from wyseos.mate import Client
from wyseos.mate.config import load_config
from wyseos.mate.models import CreateSessionRequest
from wyseos.mate.websocket import WebSocketClient, TaskExecutionOptions

# Initialize client
client = Client(load_config("mate.yaml"))

# Create session
session = client.session.create(
    CreateSessionRequest(team_id="wyse_mate", task="Analyze market trends")
)
session_info = client.session.get_info(session.session_id)

# Execute AI task
ws_client = WebSocketClient(
    base_url=client.base_url,
    api_key=client.api_key,
    session_id=session_info.session_id
)
task_runner = ws_client.create_task_runner(client, session_info)

result = task_runner.run_task(
    task="Analyze Q4 2024 market trends in tech sector",
    team_id="wyse_mate",
    options=TaskExecutionOptions(auto_accept_plan=True)
)

if result.success:
    print(f"✅ Analysis complete: {result.final_answer}")
    print(f"⏱️ Completed in {result.session_duration:.1f} seconds")
else:
    print(f"❌ Task failed: {result.error}")

# Automated execution - fire and forget
result = task_runner.run_task(
    task="Create a comprehensive market analysis report",
    team_id="wyse_mate",
    options=TaskExecutionOptions(auto_accept_plan=True)
)

# Interactive execution - with user input
task_runner.run_interactive_session(
    initial_task="Help me research competitors",
    team_id="wyse_mate"
)

# Upload files for AI analysis
uploaded_files = []
upload_result = client.file_upload.upload_file("data.csv")
if upload_result.get("file_url"):
    uploaded_files.append({
        "file_name": "data.csv",
        "file_url": upload_result["file_url"]
    })

# Use files in task execution
result = task_runner.run_task(
    task="Analyze this dataset and create visualizations",
    attachments=uploaded_files,
    team_id="wyse_mate"
)

from wyseos.mate.websocket import WebSocketClient, MessageType

ws = WebSocketClient(
    base_url=client.base_url,
    api_key=client.api_key,
    session_id=session_id
)

ws.set_message_handler(lambda msg: print(f"AI Agent: {msg.get('content')}"))
ws.connect(session_id)

# Configuration options
options = TaskExecutionOptions(
    auto_accept_plan=True,           # ✅ Auto-approve AI plans
    capture_screenshots=False,        # 📸 Browser screenshots
    enable_browser_logging=True,      # 🌐 Browser activity logs
    completion_timeout=300,           # ⏱️ Task timeout
)

# Client operations
client.user.list_api_keys()
client.team.get_list()
client.agent.get_list()
client.session.create()

# File operations
client.file_upload.validate_file()
client.file_upload.upload_file()

# Task execution
task_runner.run_task()              # Automated execution
task_runner.run_interactive_session()  # Interactive mode

from wyseos.mate.errors import APIError, ValidationError, NetworkError

try:
    result = task_runner.run_task("Your task")
except ValidationError as e:
    print(f"Validation error: {e}")
except APIError as e:
    print(f"API error: {e.message}")
except NetworkError as e:
    print(f"Network error: {e}")