# SPDX-License-Identifier: BSD-3-Clause

"""Render a rollout trace as a single HTML page you can scroll through.

A JSONL trace answers "what was the reward"; it does not show you *why*. This
lays each episode out as a filmstrip: the view the model was looking at, what it
said, the action it chose, what the environment replied, and the running cost —
ending with the reveal and the score.

Usage:
    python scripts/render_trace.py rollouts/anthropic_agentic
    python scripts/render_trace.py rollouts/hf_agentic --out qwen.html
"""

from __future__ import annotations

import argparse
import base64
import html
import json
import pathlib


def _data_uri(path: pathlib.Path) -> str:
    """Inline an image so the report is a single portable file."""
    if not path.exists():
        return ""
    mime = "image/png" if path.suffix == ".png" else "image/jpeg"
    return f"data:{mime};base64,{base64.b64encode(path.read_bytes()).decode()}"


def _fmt_distance(km: float | None) -> str:
    if km is None:
        return "—"
    return f"{km * 1000:.0f} m" if km < 10 else f"{km:,.0f} km"


def _verdict(km: float | None) -> str:
    if km is None:
        return "no guess"
    if km < 0.025:
        return "perfect"
    if km < 25:
        return "pinpoint"
    if km < 200:
        return "close"
    if km < 1500:
        return "right region"
    return "wrong continent"


def render(trace_dir: pathlib.Path, out: pathlib.Path) -> None:
    """Write one HTML report for a trace directory."""
    records = [
        json.loads(line)
        for line in (trace_dir / "trace.jsonl").read_text().splitlines()
        if line.strip()
    ]
    summary_path = trace_dir / "summary.json"
    summary = json.loads(summary_path.read_text()) if summary_path.exists() else {}

    episodes = []
    for record in records:
        turns = []
        for row in record.get("trace", []):
            image = ""
            if row.get("image"):
                image = _data_uri(trace_dir / row["image"])
            action = row.get("action") or {}
            kind = (
                action.get("action", "reset")
                if isinstance(action, dict)
                else str(action)
            )
            args = (
                ", ".join(
                    f"{key}={value}" for key, value in action.items() if key != "action"
                )
                if isinstance(action, dict)
                else ""
            )
            turns.append(
                {
                    "turn": row.get("turn"),
                    "kind": kind,
                    "args": args,
                    "reply": row.get("reply") or "",
                    "feedback": row.get("feedback") or "",
                    "steps": row.get("steps_remaining"),
                    "cost": row.get("action_cost"),
                    "image": image,
                    "image_kind": row.get("image_kind", "none"),
                }
            )
        episodes.append({"record": record, "turns": turns})

    parts = [
        _HEAD.replace(
            "__TITLE__", html.escape(str(summary.get("model", trace_dir.name)))
        )
    ]
    parts.append(_summary_block(summary, records))
    for episode in episodes:
        parts.append(_episode_block(episode))
    parts.append("</body></html>")
    out.write_text("".join(parts))
    print(f"wrote {out} ({out.stat().st_size / 1e6:.1f} MB, {len(episodes)} episodes)")


def _summary_block(summary: dict, records: list[dict]) -> str:
    rows = "".join(
        f"<div class='stat'><span>{html.escape(label)}</span><b>{html.escape(str(value))}</b></div>"
        for label, value in [
            ("model", summary.get("model", "—")),
            ("mode", summary.get("mode", "—")),
            ("episodes", summary.get("episodes", len(records))),
            ("mean reward", f"{summary.get('mean_reward', 0):.3f}"),
            ("median distance", _fmt_distance(summary.get("median_distance_km"))),
            ("within 200 km", f"{summary.get('within_200km', 0)}/{len(records)}"),
            ("parsed", f"{summary.get('parsed', 0)}/{len(records)}"),
            ("wall clock", f"{summary.get('wall_clock_s', 0):.1f}s"),
        ]
    )
    return f"<header><h1>rollout trace</h1><div class='stats'>{rows}</div></header>"


def _episode_block(episode: dict) -> str:
    record = episode["record"]
    km = record.get("distance_km")
    reward = record.get("reward") or 0.0
    head = (
        f"<h2>episode {record.get('task_index')} "
        f"<span class='country'>{html.escape(str(record.get('country', '')))}</span>"
        f"<span class='verdict v-{_verdict(km).replace(' ', '-')}'>{_verdict(km)}</span>"
        f"<span class='num'>{_fmt_distance(km)}</span>"
        f"<span class='num'>reward {reward:.3f}</span>"
        f"<span class='num'>{record.get('turns', 0)} turns</span>"
        f"<span class='num'>{record.get('latency_s', 0):.1f}s</span></h2>"
    )
    cards = []
    for turn in episode["turns"]:
        image = (
            f"<img src='{turn['image']}' alt='observation at turn {turn['turn']}'>"
            if turn["image"]
            else "<div class='noimg'>no image</div>"
        )
        reply = html.escape(turn["reply"][:1400])
        meta = []
        if turn["steps"] is not None:
            meta.append(f"{turn['steps']} steps left")
        if turn["cost"] is not None:
            meta.append(f"cost {turn['cost']:.2f}")
        cards.append(
            f"""<div class='turn {"map" if turn["image_kind"] == "map" else ""}'>
  <div class='thumb'>{image}</div>
  <div class='body'>
    <div class='act'><b>{html.escape(str(turn["kind"]))}</b>
      <span>{html.escape(turn["args"])}</span>
      <span class='meta'>{html.escape(" · ".join(meta))}</span></div>
    <div class='fb'>{html.escape(turn["feedback"])}</div>
    {f"<details><summary>model reply</summary><pre>{reply}</pre></details>" if reply else ""}
  </div>
</div>"""
        )
    return f"<section>{head}<div class='turns'>{''.join(cards)}</div></section>"


_HEAD = """<!doctype html>
<html lang="en"><head><meta charset="utf-8">
<meta name="viewport" content="width=device-width,initial-scale=1">
<title>__TITLE__ — rollout trace</title>
<style>
  :root {
    --ground:#eef1f4; --panel:#fff; --edge:#d3dae0; --ink:#1d2328;
    --soft:#5c666e; --faint:#8b959c; --wire:#4a7d99; --pin:#c4332a; --good:#3f7a55;
  }
  @media (prefers-color-scheme:dark){ :root{
    --ground:#12171b; --panel:#191f24; --edge:#2b343b; --ink:#e6ebef;
    --soft:#a3aeb6; --faint:#6e7a83; --wire:#7fb3cc; --pin:#e8695e; --good:#79b892; } }
  *{box-sizing:border-box}
  body{margin:0;background:var(--ground);color:var(--ink);
    font:14px/1.6 ui-monospace,"SF Mono",Menlo,monospace;padding:24px 18px 60px}
  header{max-width:1180px;margin:0 auto 26px}
  h1{font-size:19px;margin:0 0 12px;font-weight:600;letter-spacing:.02em}
  .stats{display:flex;flex-wrap:wrap;gap:18px}
  .stat{background:var(--panel);border:1px solid var(--edge);border-radius:4px;
    padding:7px 12px;min-width:118px}
  .stat span{display:block;font-size:9.5px;letter-spacing:.12em;
    text-transform:uppercase;color:var(--faint)}
  .stat b{font-size:15px;font-weight:500;font-variant-numeric:tabular-nums}
  section{max-width:1180px;margin:0 auto 30px;background:var(--panel);
    border:1px solid var(--edge);border-radius:5px;overflow:hidden}
  h2{margin:0;padding:11px 15px;font-size:13px;font-weight:600;
    border-bottom:1px solid var(--edge);display:flex;align-items:center;
    gap:12px;flex-wrap:wrap}
  h2 .country{color:var(--soft);font-weight:400}
  h2 .num{margin-left:auto;font-weight:400;color:var(--soft);font-size:11.5px;
    font-variant-numeric:tabular-nums}
  h2 .num + .num{margin-left:0}
  .verdict{font-size:10px;letter-spacing:.1em;text-transform:uppercase;
    padding:2px 7px;border-radius:3px;border:1px solid var(--edge);color:var(--soft)}
  .v-pinpoint,.v-perfect{color:var(--good);border-color:var(--good)}
  .v-wrong-continent,.v-no-guess{color:var(--pin);border-color:var(--pin)}
  .turns{display:flex;flex-direction:column}
  .turn{display:flex;gap:14px;padding:12px 15px;border-bottom:1px solid var(--edge)}
  .turn:last-child{border-bottom:none}
  .turn.map{background:color-mix(in srgb,var(--wire) 7%,transparent)}
  .thumb{flex:0 0 200px}
  .thumb img{width:200px;border-radius:3px;display:block;border:1px solid var(--edge)}
  .noimg{width:200px;height:96px;border:1px dashed var(--edge);border-radius:3px;
    display:flex;align-items:center;justify-content:center;color:var(--faint);font-size:11px}
  .body{flex:1;min-width:0}
  .act b{color:var(--wire)}
  .act span{color:var(--soft);font-size:12px;margin-left:6px}
  .act .meta{float:right;color:var(--faint);font-size:11px}
  .fb{margin-top:5px;color:var(--soft);font-size:12.5px}
  details{margin-top:8px}
  summary{cursor:pointer;font-size:11px;color:var(--faint)}
  pre{white-space:pre-wrap;font-size:11.5px;background:var(--ground);
    border:1px solid var(--edge);border-radius:3px;padding:9px;margin:7px 0 0;
    max-height:260px;overflow:auto}
</style></head><body>
"""


def main() -> None:
    """Command-line entry point."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("trace_dir", type=pathlib.Path)
    parser.add_argument("--out", type=pathlib.Path)
    args = parser.parse_args()
    out = args.out or (args.trace_dir / "trace.html")
    render(args.trace_dir, out)


if __name__ == "__main__":
    main()