#!/usr/bin/env python3
"""
Run a CUA task remotely against cua.derlem.com (or any CUA instance).

Connects via WebSocket, sends a prompt, streams progress to the terminal,
and saves the full trace as JSON.

Usage:
  poetry install
  poetry run cua-run "Open Firefox and go to example.com"

Environment (optional overrides):
  CUA_BASE_URL   Base URL (default: https://cua.derlem.com)
  CUA_USER       HTTP Basic Auth username (default: cua)
  CUA_PASSWORD   HTTP Basic Auth password (if your CUA instance is protected)
"""

from __future__ import annotations

import argparse
import asyncio
import base64
import json
import os
import sys
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen

try:
    import websockets
    from websockets.exceptions import ConnectionClosed
except ImportError:
    print("Missing dependency: poetry install", file=sys.stderr)
    sys.exit(1)

from sicss2026.html_exporter import save_trace_html

DEFAULT_BASE_URL = "https://cua.derlem.com"
DEFAULT_BASIC_AUTH_USER = "cua"
# 8B is listed in CUA but often unavailable on enabled HF providers; 30B works via Novita.
DEFAULT_MODEL = "Qwen/Qwen3-VL-30B-A3B-Instruct"
DEFAULT_MAX_STEPS = 200
DEFAULT_USER_AGENT = "sicss2026-cua-run/0.1"

HF_PROVIDERS_URL = "https://huggingface.co/settings/inference-providers"
HF_MODELS_URL = "https://huggingface.co/inference/models"


class TaskFailure(RuntimeError):
    """CUA task failed; carries user-facing hints beyond the server message."""

    def __init__(
        self,
        message: str,
        *,
        model_id: str,
        final_state: str | None = None,
        server_message: str | None = None,
        hints: list[str] | None = None,
        result: dict[str, Any] | None = None,
    ) -> None:
        super().__init__(message)
        self.model_id = model_id
        self.final_state = final_state
        self.server_message = server_message
        self.hints = hints or []
        self.result = result


def fetch_available_models(base_url: str, headers: dict[str, str]) -> list[str]:
    req = Request(
        f"{base_url.rstrip('/')}/api/models",
        headers=headers,
    )
    with urlopen(req, timeout=30) as resp:
        data = json.loads(resp.read().decode())
    return list(data.get("models") or [])


def failure_hints(
    *,
    model_id: str,
    server_message: str | None,
    final_state: str | None,
    available_models: list[str] | None = None,
) -> list[str]:
    hints: list[str] = []
    generic = server_message in {None, "", "Error processing task"}

    if generic or (server_message and "model_not_supported" in server_message):
        hints.append(
            f"The model {model_id!r} may not be routed to any HF Inference Provider "
            f"enabled on the server account."
        )
        hints.append(
            f"Open Inference Providers settings and confirm a provider serves this model: "
            f"{HF_PROVIDERS_URL}"
        )
        if available_models:
            other = [m for m in available_models if m != model_id]
            if other:
                hints.append(
                    "Try another model listed by the server, e.g. "
                    + " or ".join(f"--model {m}" for m in other[:2])
                )
        hints.append(
            "Workshop default that worked on Novita: "
            "--model Qwen/Qwen3-VL-30B-A3B-Instruct"
        )
        hints.append(f"Browse provider coverage: {HF_MODELS_URL}")

    if final_state == "max_sandboxes_reached":
        hints.append("E2B sandbox pool is full. Wait and retry, or stop other sessions.")

    if final_state == "sandbox_timeout":
        hints.append("The E2B sandbox timed out. Try a shorter or simpler task.")

    if final_state == "max_steps_reached":
        hints.append("The agent hit its step limit before finishing. Simplify the prompt.")

    if generic:
        hints.append(
            "The server hides the root cause in WebSocket messages. "
            "Check container logs on the host: "
            "docker compose -f /home/artdata/cua/docker-compose.yml logs cua --tail 50"
        )

    return hints


def format_failure(exc: TaskFailure) -> str:
    lines = [str(exc)]
    if exc.server_message and exc.server_message != str(exc):
        lines.append(f"Server message: {exc.server_message}")
    if exc.final_state:
        lines.append(f"Final state: {exc.final_state}")
    lines.append(f"Model: {exc.model_id}")
    if exc.hints:
        lines.append("")
        lines.append("Likely fixes:")
        for hint in exc.hints:
            lines.append(f"  - {hint}")
    return "\n".join(lines)


def build_task_failure(
    *,
    model_id: str,
    server_message: str | None = None,
    final_state: str | None = None,
    available_models: list[str] | None = None,
    result: dict[str, Any] | None = None,
) -> TaskFailure:
    if final_state == "max_sandboxes_reached":
        headline = "Server at sandbox capacity (max_sandboxes_reached)"
    elif final_state == "sandbox_timeout":
        headline = "E2B sandbox timed out"
    elif final_state == "max_steps_reached":
        headline = "Agent reached max steps without completing the task"
    elif final_state == "stopped":
        headline = "Task was stopped"
    elif server_message and server_message != "Error processing task":
        headline = server_message
    elif final_state == "error":
        headline = "Agent run failed on the server"
    else:
        headline = server_message or "Task failed"

    hints = failure_hints(
        model_id=model_id,
        server_message=server_message,
        final_state=final_state,
        available_models=available_models,
    )
    return TaskFailure(
        headline,
        model_id=model_id,
        final_state=final_state,
        server_message=server_message,
        hints=hints,
        result=result,
    )


def repo_root() -> Path:
    return Path(__file__).resolve().parent.parent


def default_output_dir() -> Path:
    return repo_root() / "scripts" / "output"


def auth_headers(user: str, password: str) -> dict[str, str]:
    headers = {"User-Agent": DEFAULT_USER_AGENT}
    if password:
        headers["Authorization"] = basic_auth_header(user, password)
    return headers


def basic_auth_header(user: str, password: str) -> str:
    token = base64.b64encode(f"{user}:{password}".encode()).decode("ascii")
    return f"Basic {token}"


def http_to_ws_url(base_url: str) -> str:
    base = base_url.rstrip("/")
    if base.startswith("https://"):
        return "wss://" + base[len("https://") :] + "/ws"
    if base.startswith("http://"):
        return "ws://" + base[len("http://") :] + "/ws"
    return "wss://" + base + "/ws"


def utc_now_iso() -> str:
    return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")


def extract_final_answer(steps: list[dict[str, Any]]) -> str | None:
    for step in reversed(steps):
        for action in step.get("actions") or []:
            if isinstance(action, dict) and action.get("function_name") == "final_answer":
                params = action.get("parameters") or {}
                answer = params.get("answer") or params.get("arg_0")
                if answer is not None:
                    return str(answer)
        for action in step.get("actions") or []:
            if isinstance(action, str) and action.startswith("Final answer:"):
                return action.removeprefix("Final answer:").strip()
    if steps:
        thought = steps[-1].get("thought")
        if thought:
            return str(thought)
    return None


def summarize_action(action: dict[str, Any] | str) -> str:
    if isinstance(action, str):
        return action
    name = action.get("function_name", "unknown")
    params = action.get("parameters") or {}
    if name == "final_answer":
        return f"final_answer: {params.get('answer') or params.get('arg_0')}"
    return name


def check_health(base_url: str, headers: dict[str, str]) -> dict[str, Any]:
    req = Request(
        f"{base_url.rstrip('/')}/health",
        headers=headers,
    )
    with urlopen(req, timeout=30) as resp:
        return json.loads(resp.read().decode())


def build_user_task(
    trace_id: str,
    instruction: str,
    model_id: str,
    *,
    max_steps: int = DEFAULT_MAX_STEPS,
) -> dict[str, Any]:
    return {
        "type": "user_task",
        "trace": {
            "id": trace_id,
            "instruction": instruction,
            "modelId": model_id,
            "timestamp": utc_now_iso(),
            "isRunning": True,
            "steps": [],
            "traceMetadata": {
                "traceId": trace_id,
                "inputTokensUsed": 0,
                "outputTokensUsed": 0,
                "duration": 0.0,
                "numberOfSteps": 0,
                "maxSteps": max_steps,
                "completed": False,
                "final_state": None,
                "user_evaluation": "not_evaluated",
            },
        },
    }


async def run_task(
    *,
    ws_url: str,
    headers: dict[str, str],
    instruction: str,
    model_id: str,
    timeout_seconds: float,
    include_images: bool,
    available_models: list[str] | None = None,
    max_steps: int = DEFAULT_MAX_STEPS,
) -> dict[str, Any]:
    result: dict[str, Any] = {
        "instruction": instruction,
        "modelId": model_id,
        "startedAt": utc_now_iso(),
        "traceId": None,
        "vncUrl": None,
        "steps": [],
        "completion": None,
        "error": None,
        "failed": False,
    }
    server_error: str | None = None

    async with websockets.connect(
        ws_url,
        extra_headers=headers,
        open_timeout=30,
        ping_interval=20,
        ping_timeout=20,
        max_size=50 * 1024 * 1024,
    ) as ws:
        trace_id: str | None = None
        deadline = asyncio.get_event_loop().time() + timeout_seconds

        while True:
            remaining = deadline - asyncio.get_event_loop().time()
            if remaining <= 0:
                raise TimeoutError(f"No completion within {timeout_seconds:.0f}s")

            raw = await asyncio.wait_for(ws.recv(), timeout=remaining)
            event = json.loads(raw)
            event_type = event.get("type")

            if event_type == "heartbeat":
                trace_id = event["uuid"]
                result["traceId"] = trace_id
                print(f"Session ready (trace id: {trace_id})")
                await ws.send(
                    json.dumps(
                        build_user_task(trace_id, instruction, model_id, max_steps=max_steps)
                    )
                )
                print(f"Task sent: {instruction!r}")
                continue

            if event_type == "vnc_url_set":
                result["vncUrl"] = event.get("vncUrl")
                print(f"VNC stream: {result['vncUrl']}")
                continue

            if event_type == "agent_start":
                status = event.get("status")
                if status == "max_sandboxes_reached":
                    raise build_task_failure(
                        model_id=model_id,
                        final_state="max_sandboxes_reached",
                        available_models=available_models,
                        result=result,
                    )
                print("Agent started")
                continue

            if event_type == "agent_progress":
                step = event.get("agentStep") or {}
                step_record = dict(step)
                if not include_images and "image" in step_record:
                    step_record["image"] = "<omitted>"
                result["steps"].append(step_record)
                result["traceMetadata"] = event.get("traceMetadata")

                step_id = step.get("stepId", "?")
                thought = (step.get("thought") or "").strip()
                actions = step.get("actions") or []
                action_text = "; ".join(summarize_action(a) for a in actions[:3])
                print(f"  step {step_id}: {thought[:120]}{'…' if len(thought) > 120 else ''}")
                if action_text:
                    print(f"           actions: {action_text}")
                if step.get("error"):
                    print(f"           step error: {step['error']}")
                continue

            if event_type == "agent_complete":
                metadata = event.get("traceMetadata") or {}
                final_state = event.get("final_state")
                result["traceMetadata"] = metadata
                result["completion"] = {
                    "final_state": final_state,
                    "finalAnswer": extract_final_answer(result["steps"]),
                    "numberOfSteps": metadata.get("numberOfSteps"),
                    "duration": metadata.get("duration"),
                    "inputTokensUsed": metadata.get("inputTokensUsed"),
                    "outputTokensUsed": metadata.get("outputTokensUsed"),
                }
                print(f"Done: {final_state}")
                answer = result["completion"]["finalAnswer"]
                if answer:
                    print(f"Final answer: {answer}")

                if final_state != "success":
                    result["failed"] = True
                    result["error"] = server_error or f"final_state={final_state}"
                    result["finishedAt"] = utc_now_iso()
                    raise build_task_failure(
                        model_id=model_id,
                        server_message=server_error,
                        final_state=final_state,
                        available_models=available_models,
                        result=result,
                    )
                break

            if event_type == "agent_error":
                server_error = event.get("error", "unknown error")
                result["error"] = server_error
                # Server also sends agent_complete; wait for it unless the socket ends.
                print(f"Server error event: {server_error}")
                continue

            if event_type in {"vnc_url_unset"}:
                continue

            print(f"  (ignored event: {event_type})")

    result["finishedAt"] = utc_now_iso()
    if server_error:
        result["failed"] = True
        raise build_task_failure(
            model_id=model_id,
            server_message=server_error,
            available_models=available_models,
            result=result,
        )
    return result


def save_trace(
    result: dict[str, Any],
    output_path: Path | None,
    *,
    write_html: bool = True,
) -> tuple[Path, Path | None]:
    if output_path is None:
        stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
        output_path = default_output_dir() / f"trace-{stamp}.json"
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(json.dumps(result, indent=2), encoding="utf-8")

    html_path: Path | None = None
    if write_html:
        html_path = output_path.with_suffix(".html")
        save_trace_html(result, html_path)
    return output_path, html_path


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Run a CUA task remotely via WebSocket and save the trace.",
    )
    parser.add_argument(
        "instruction",
        help="Natural-language task for the computer-use agent",
    )
    parser.add_argument(
        "--base-url",
        default=os.environ.get("CUA_BASE_URL", DEFAULT_BASE_URL),
        help=f"CUA base URL (default: {DEFAULT_BASE_URL})",
    )
    parser.add_argument(
        "--user",
        default=os.environ.get("CUA_USER", DEFAULT_BASIC_AUTH_USER),
        help=f"Basic auth username (default: {DEFAULT_BASIC_AUTH_USER})",
    )
    parser.add_argument(
        "--password",
        default=os.environ.get("CUA_PASSWORD", ""),
        help="Basic auth password (optional; or set CUA_PASSWORD)",
    )
    parser.add_argument(
        "--model",
        default=DEFAULT_MODEL,
        help=f"Vision model id (default: {DEFAULT_MODEL})",
    )
    parser.add_argument(
        "--output",
        "-o",
        type=Path,
        help="Output JSON path (default: scripts/output/trace-<timestamp>.json)",
    )
    parser.add_argument(
        "--timeout",
        type=float,
        default=1800,
        help="Max seconds to wait for completion (default: 1800)",
    )
    parser.add_argument(
        "--no-images",
        action="store_true",
        help="Omit base64 screenshots from saved JSON",
    )
    parser.add_argument(
        "--no-html",
        action="store_true",
        help="Do not write an HTML report alongside the JSON trace",
    )
    parser.add_argument(
        "--skip-health-check",
        action="store_true",
        help="Skip GET /health before connecting",
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()

    request_headers = auth_headers(args.user, args.password)
    base_url = args.base_url.rstrip("/")
    ws_url = http_to_ws_url(base_url)

    if not args.skip_health_check:
        try:
            health = check_health(base_url, request_headers)
            print(f"Health: {health.get('status', health)}")
        except HTTPError as exc:
            print(f"Health check failed: HTTP {exc.code}", file=sys.stderr)
            return 1
        except URLError as exc:
            print(f"Health check failed: {exc.reason}", file=sys.stderr)
            return 1

    print(f"Connecting to {ws_url}")

    available_models: list[str] | None = None
    try:
        available_models = fetch_available_models(base_url, request_headers)
        print(f"Models: {', '.join(available_models)}")
    except (HTTPError, URLError, json.JSONDecodeError, KeyError):
        pass

    result: dict[str, Any] | None = None
    try:
        result = asyncio.run(
            run_task(
                ws_url=ws_url,
                headers=request_headers,
                instruction=args.instruction,
                model_id=args.model,
                timeout_seconds=args.timeout,
                include_images=not args.no_images,
                available_models=available_models,
            )
        )
    except TimeoutError as exc:
        print(f"Error: {exc}", file=sys.stderr)
        return 1
    except ConnectionClosed as exc:
        print(f"WebSocket closed unexpectedly: {exc}", file=sys.stderr)
        return 1
    except TaskFailure as exc:
        partial = exc.result
        if partial is not None:
            json_path, html_path = save_trace(
                partial, args.output, write_html=not args.no_html
            )
            print(f"Partial trace saved: {json_path}")
            if html_path:
                print(f"HTML report saved: {html_path}")
        print(format_failure(exc), file=sys.stderr)
        return 1
    except RuntimeError as exc:
        print(f"Error: {exc}", file=sys.stderr)
        return 1

    json_path, html_path = save_trace(result, args.output, write_html=not args.no_html)
    print(f"Trace saved: {json_path}")
    if html_path:
        print(f"HTML report saved: {html_path}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
