"""Export a CUA trace JSON as a self-contained HTML report (CUA web UI style)."""

from __future__ import annotations

import html
import json
from datetime import datetime
from pathlib import Path
from typing import Any


FINAL_STATE_LABELS = {
    "success": ("Task completed", "success"),
    "stopped": ("Task stopped", "warning"),
    "max_steps_reached": ("Max steps reached", "warning"),
    "sandbox_timeout": ("Sandbox timeout", "error"),
    "error": ("Task failed", "error"),
}


def _esc(value: Any) -> str:
    if value is None:
        return ""
    return html.escape(str(value))


def _action_label(action: dict[str, Any] | str) -> str:
    if isinstance(action, str):
        return action
    if action.get("description"):
        return str(action["description"])
    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 _format_duration(seconds: float | None) -> str:
    if seconds is None:
        return "—"
    return f"{seconds:.1f}s"


def export_trace_to_html(trace: dict[str, Any], *, title: str | None = None) -> str:
    instruction = trace.get("instruction", "")
    model_id = trace.get("modelId", "")
    trace_id = trace.get("traceId", "")
    started_at = trace.get("startedAt", "")
    finished_at = trace.get("finishedAt", "")
    vnc_url = trace.get("vncUrl")
    steps: list[dict[str, Any]] = trace.get("steps") or []
    completion = trace.get("completion") or {}
    metadata = trace.get("traceMetadata") or {}
    final_state = completion.get("final_state") or metadata.get("final_state")
    final_answer = completion.get("finalAnswer")
    page_title = title or f"CUA trace {trace_id or 'export'}"

    if final_state:
        final_label, final_tone = FINAL_STATE_LABELS.get(
            final_state, (f"Finished: {final_state}", "neutral")
        )
    else:
        final_label, final_tone = ("In progress", "neutral")

    total_duration = completion.get("duration") or metadata.get("duration")
    input_tokens = completion.get("inputTokensUsed") or metadata.get("inputTokensUsed", 0)
    output_tokens = completion.get("outputTokensUsed") or metadata.get("outputTokensUsed", 0)
    step_count = completion.get("numberOfSteps") or metadata.get("numberOfSteps", len(steps))
    max_steps = metadata.get("maxSteps", "—")

    step_cards: list[str] = []
    for index, step in enumerate(steps):
        step_id = step.get("stepId", index + 1)
        image = step.get("image")
        image_html = ""
        if image and image != "<omitted>":
            src = _esc(image)
            image_html = f'<img class="step-image" src="{src}" alt="Step {step_id} screenshot" loading="lazy" />'

        actions = step.get("actions") or []
        actions_html = "".join(
            f'<div class="action-line"><span class="action-arrow">→</span> {_esc(_action_label(a))}</div>'
            for a in actions
        )

        thought_html = ""
        if step.get("thought"):
            thought_html = f"""
            <details class="thought" open>
              <summary>Thought</summary>
              <pre>{_esc(step["thought"])}</pre>
            </details>
            """

        error_html = ""
        if step.get("error"):
            error_html = f'<div class="step-error">Error: {_esc(step["error"])}</div>'

        step_cards.append(
            f"""
            <article class="step-card" id="step-{index + 1}">
              <header class="step-header">
                <span class="step-badge">{index + 1}</span>
                <div class="chips">
                  <span class="chip">⏱ {_esc(_format_duration(step.get("duration")))}</span>
                  <span class="chip">in {int(step.get("inputTokensUsed") or 0):,}</span>
                  <span class="chip">out {int(step.get("outputTokensUsed") or 0):,}</span>
                </div>
              </header>
              {image_html}
              {f'<section class="actions"><div class="section-label">Action</div>{actions_html}</section>' if actions_html else ""}
              {thought_html}
              {error_html}
            </article>
            """
        )

    timeline_dots = "".join(
        f'<span class="timeline-dot{" active" if i == len(steps) - 1 else " done" if i < len(steps) else ""}">{i + 1}</span>'
        for i in range(max(len(steps), 1))
    )

    final_answer_html = ""
    if final_answer:
        final_answer_html = f"""
        <section class="final-answer">
          <h2>Final answer</h2>
          <pre>{_esc(final_answer)}</pre>
        </section>
        """

    vnc_html = ""
    if vnc_url:
        vnc_html = f'<p class="meta-line">VNC replay: <a href="{_esc(vnc_url)}" target="_blank" rel="noopener">{_esc(vnc_url)}</a></p>'

    exported_at = datetime.now().astimezone().isoformat(timespec="seconds")

    return f"""<!DOCTYPE html>
<html lang="en">
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1" />
  <title>{_esc(page_title)}</title>
  <style>
    :root {{
      --primary: #4F86C6;
      --bg: #f8f9fa;
      --paper: #ffffff;
      --text: rgba(0, 0, 0, 0.87);
      --text-secondary: rgba(0, 0, 0, 0.6);
      --border: rgba(0, 0, 0, 0.15);
      --success: #2e7d32;
      --warning: #ed6c02;
      --error: #d32f2f;
      --radius: 10px;
      --card-radius: 12px;
      --font: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, "Helvetica Neue", Arial, sans-serif;
    }}
    * {{ box-sizing: border-box; }}
    body {{
      margin: 0;
      font-family: var(--font);
      background: var(--bg);
      color: var(--text);
      line-height: 1.5;
    }}
    .page {{
      max-width: 960px;
      margin: 0 auto;
      padding: 24px 16px 48px;
    }}
    .hero, .step-card, .final-card, .final-answer {{
      background: var(--paper);
      border: 1px solid var(--border);
      border-radius: var(--card-radius);
      box-shadow: 0 1px 2px rgba(0,0,0,0.04);
    }}
    .hero {{ padding: 20px; margin-bottom: 16px; }}
    .hero h1 {{
      margin: 0 0 8px;
      font-size: 1.25rem;
      color: var(--primary);
    }}
    .instruction {{
      font-size: 1rem;
      margin: 0 0 12px;
      white-space: pre-wrap;
    }}
    .meta-line {{
      margin: 4px 0;
      color: var(--text-secondary);
      font-size: 0.875rem;
      word-break: break-word;
    }}
    .summary-chips {{
      display: flex;
      flex-wrap: wrap;
      gap: 8px;
      margin-top: 12px;
    }}
    .summary-chips .chip {{
      background: rgba(0,0,0,0.04);
      border-radius: 6px;
      padding: 4px 10px;
      font-size: 0.8rem;
      font-weight: 600;
    }}
    .timeline {{
      display: flex;
      gap: 10px;
      overflow-x: auto;
      padding: 12px 4px 16px;
      margin-bottom: 8px;
    }}
    .timeline-dot {{
      width: 36px;
      height: 36px;
      border-radius: 50%;
      display: inline-flex;
      align-items: center;
      justify-content: center;
      font-size: 0.8rem;
      font-weight: 700;
      border: 2px solid #cfd8dc;
      color: var(--text-secondary);
      flex: 0 0 auto;
      background: #fff;
    }}
    .timeline-dot.done {{
      background: var(--primary);
      border-color: var(--primary);
      color: #fff;
    }}
    .timeline-dot.active {{
      background: #fff;
      border-color: var(--primary);
      color: var(--primary);
      box-shadow: 0 0 0 4px rgba(79, 134, 198, 0.15);
    }}
    .steps-title {{
      font-size: 0.95rem;
      font-weight: 700;
      margin: 8px 0 12px;
      color: var(--text-secondary);
      text-transform: uppercase;
      letter-spacing: 0.04em;
    }}
    .step-card {{
      padding: 14px;
      margin-bottom: 12px;
    }}
    .step-header {{
      display: flex;
      align-items: center;
      justify-content: space-between;
      gap: 12px;
      margin-bottom: 10px;
    }}
    .step-badge {{
      width: 28px;
      height: 28px;
      border-radius: 50%;
      background: var(--primary);
      color: #fff;
      display: inline-flex;
      align-items: center;
      justify-content: center;
      font-weight: 700;
      font-size: 0.85rem;
      flex: 0 0 auto;
    }}
    .chips {{
      display: flex;
      flex-wrap: wrap;
      gap: 6px;
      justify-content: flex-end;
    }}
    .chip {{
      background: rgba(0,0,0,0.05);
      border-radius: 6px;
      padding: 2px 8px;
      font-size: 0.7rem;
      font-weight: 600;
      color: var(--text);
    }}
    .step-image {{
      width: 100%;
      border-radius: 8px;
      border: 1px solid var(--border);
      background: rgba(0,0,0,0.03);
      margin-bottom: 10px;
    }}
    .section-label {{
      font-size: 0.75rem;
      font-weight: 700;
      color: var(--text-secondary);
      margin-bottom: 4px;
      text-transform: uppercase;
      letter-spacing: 0.03em;
    }}
    .action-line {{
      font-family: ui-monospace, SFMono-Regular, Menlo, Monaco, Consolas, monospace;
      font-size: 0.85rem;
      padding: 4px 0;
    }}
    .action-arrow {{
      color: var(--primary);
      margin-right: 6px;
    }}
    .thought {{
      margin-top: 8px;
      border-top: 1px solid var(--border);
      padding-top: 8px;
    }}
    .thought summary {{
      cursor: pointer;
      font-weight: 600;
      color: var(--text-secondary);
    }}
    .thought pre, .final-answer pre {{
      white-space: pre-wrap;
      word-break: break-word;
      margin: 8px 0 0;
      font: inherit;
      color: var(--text);
    }}
    .step-error {{
      margin-top: 8px;
      padding: 8px 10px;
      border-radius: 8px;
      background: rgba(211, 47, 47, 0.08);
      border: 1px solid rgba(211, 47, 47, 0.35);
      color: var(--error);
      font-size: 0.875rem;
    }}
    .final-card {{
      padding: 16px;
      margin: 16px 0;
      border-width: 2px;
    }}
    .final-card.success {{ border-color: rgba(46, 125, 50, 0.5); }}
    .final-card.warning {{ border-color: rgba(237, 108, 2, 0.5); }}
    .final-card.error {{ border-color: rgba(211, 47, 47, 0.5); }}
    .final-card.neutral {{ border-color: var(--border); }}
    .final-card h2 {{
      margin: 0;
      font-size: 1rem;
    }}
    .final-answer {{
      padding: 16px;
      margin-top: 16px;
    }}
    .final-answer h2 {{
      margin: 0 0 8px;
      font-size: 1rem;
      color: var(--primary);
    }}
    footer {{
      margin-top: 24px;
      color: var(--text-secondary);
      font-size: 0.8rem;
    }}
    a {{ color: var(--primary); }}
  </style>
</head>
<body>
  <div class="page">
    <section class="hero">
      <h1>Computer Use Agent — trace</h1>
      <p class="instruction">{_esc(instruction)}</p>
      <p class="meta-line"><strong>Model:</strong> {_esc(model_id)}</p>
      <p class="meta-line"><strong>Trace ID:</strong> {_esc(trace_id)}</p>
      <p class="meta-line"><strong>Started:</strong> {_esc(started_at)}</p>
      <p class="meta-line"><strong>Finished:</strong> {_esc(finished_at)}</p>
      {vnc_html}
      <div class="summary-chips">
        <span class="chip">Steps {step_count}/{_esc(max_steps)}</span>
        <span class="chip">Duration {_esc(_format_duration(total_duration))}</span>
        <span class="chip">Tokens in {int(input_tokens or 0):,}</span>
        <span class="chip">Tokens out {int(output_tokens or 0):,}</span>
      </div>
    </section>

    <div class="timeline" aria-label="Step timeline">{timeline_dots}</div>

    <div class="steps-title">Steps ({len(steps)})</div>
    {"".join(step_cards) if step_cards else '<p class="meta-line">No steps recorded.</p>'}

    <section class="final-card {final_tone}">
      <h2>{_esc(final_label)}</h2>
    </section>

    {final_answer_html}

    <footer>Exported by sicss2026 cua-run at {_esc(exported_at)}</footer>
  </div>
</body>
</html>
"""


def save_trace_html(trace: dict[str, Any], output_path: Path, *, title: str | None = None) -> Path:
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(export_trace_to_html(trace, title=title), encoding="utf-8")
    return output_path


def json_file_to_html(json_path: Path, output_path: Path | None = None) -> Path:
    trace = json.loads(json_path.read_text(encoding="utf-8"))
    if output_path is None:
        output_path = json_path.with_suffix(".html")
    return save_trace_html(trace, output_path)
