"""Run CUA against a single mock-shop scenario page."""

from __future__ import annotations

import argparse
import asyncio
import json
import sys
from pathlib import Path

from sicss2026.kit_paths import default_persona_prompts
from sicss2026.prediction_parser import parse_prediction
from sicss2026.run_remote_task import (
    DEFAULT_BASE_URL,
    DEFAULT_BASIC_AUTH_PASSWORD,
    DEFAULT_BASIC_AUTH_USER,
    DEFAULT_MODEL,
    auth_headers,
    check_health,
    fetch_available_models,
    format_failure,
    http_to_ws_url,
    run_task,
    save_trace,
    TaskFailure,
)
from urllib.error import HTTPError, URLError
from websockets.exceptions import ConnectionClosed

DEFAULT_SHOP_URL = "https://sicss2026.derlem.com"
DEFAULT_MAX_STEPS = 8

SCENARIO_PROMPT = """Open Firefox and go to {page_url}
You are acting as this shopper persona:

{persona_prompt}

Inspect ONLY what is visible on the page. Do not click checkout or add to cart.
At the end output exactly:

PURCHASE_PROBABILITY: <number between 0 and 1>
REASON: <one sentence grounded in the page>

Then stop immediately."""


def load_persona_prompt(prompts_dir: Path, persona_id: str) -> str:
    path = prompts_dir / f"{persona_id}.txt"
    if not path.exists():
        raise FileNotFoundError(f"Missing persona prompt: {path}")
    return path.read_text(encoding="utf-8").strip()


def build_instruction(
    *,
    shop_base_url: str,
    record_id: str,
    persona_prompt: str,
) -> str:
    page_url = f"{shop_base_url.rstrip('/')}/scenarios/{record_id}.html"
    return SCENARIO_PROMPT.format(page_url=page_url, persona_prompt=persona_prompt)


def extract_answer(result: dict) -> str:
    completion = result.get("completion") or {}
    if completion.get("answer"):
        return str(completion["answer"])
    steps = result.get("steps") or []
    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:
                    return str(answer)
        thought = step.get("thought")
        if thought:
            return str(thought)
    return ""


def cmd_scenario(args: argparse.Namespace) -> int:
    persona_prompt = (
        args.persona_prompt.read_text(encoding="utf-8").strip()
        if args.persona_prompt
        else load_persona_prompt(args.persona_prompts, args.persona_id)
    )
    instruction = build_instruction(
        shop_base_url=args.shop_url,
        record_id=args.record_id,
        persona_prompt=persona_prompt,
    )

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

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

    available_models = None
    try:
        available_models = fetch_available_models(base_url, headers)
    except (HTTPError, URLError, json.JSONDecodeError, KeyError):
        pass

    try:
        result = asyncio.run(
            run_task(
                ws_url=ws_url,
                headers=headers,
                instruction=instruction,
                model_id=args.model,
                timeout_seconds=args.timeout,
                include_images=not args.no_images,
                available_models=available_models,
                max_steps=args.max_steps,
            )
        )
    except TaskFailure as exc:
        print(format_failure(exc), file=sys.stderr)
        return 1
    except (TimeoutError, ConnectionClosed, RuntimeError) as exc:
        print(f"Error: {exc}", file=sys.stderr)
        return 1

    answer = extract_answer(result)
    try:
        parsed = parse_prediction(answer)
    except ValueError:
        parsed = {"purchase_probability": None, "reason": answer}

    output = {
        "record_id": args.record_id,
        "persona_id": args.persona_id,
        **parsed,
        "raw_answer": answer,
    }
    print(json.dumps(output, indent=2))

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


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Run CUA on one mock-shop scenario.")
    parser.add_argument("--record-id", required=True)
    parser.add_argument("--persona-id", required=True)
    parser.add_argument("--persona-name", required=True)
    parser.add_argument("--persona-prompts", type=Path, default=default_persona_prompts())
    parser.add_argument("--persona-prompt", type=Path)
    parser.add_argument("--shop-url", default=DEFAULT_SHOP_URL)
    parser.add_argument("--cua-url", default=DEFAULT_BASE_URL)
    parser.add_argument("--user", default=DEFAULT_BASIC_AUTH_USER)
    parser.add_argument("--password", default=DEFAULT_BASIC_AUTH_PASSWORD)
    parser.add_argument("--model", default=DEFAULT_MODEL)
    parser.add_argument("--max-steps", type=int, default=DEFAULT_MAX_STEPS)
    parser.add_argument("--timeout", type=float, default=600)
    parser.add_argument("--output", "-o", type=Path)
    parser.add_argument("--no-images", action="store_true")
    parser.add_argument("--no-html", action="store_true")
    parser.add_argument("--skip-health-check", action="store_true")
    return parser


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    try:
        return cmd_scenario(args)
    except FileNotFoundError as exc:
        print(f"Error: {exc}", file=sys.stderr)
        return 1


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