"""Build a local HTML report for Pong RAM policy experiments."""

from __future__ import annotations

import argparse
import csv
import html
import json
from collections import defaultdict
from pathlib import Path


TRIAL_FIELDS = (
    ("round", "轮次"),
    ("timestamp", "时间"),
    ("policy_name", "策略"),
    ("seed", "seed"),
    ("episodes", "局数"),
    ("score_mean", "均分"),
    ("score_min", "最低"),
    ("score_max", "最高"),
    ("env_steps", "步数"),
    ("notes", "改动"),
    ("suggestion", "建议"),
)

EPISODE_FIELDS = (
    ("episode", "局"),
    ("seed", "seed"),
    ("score", "分数"),
    ("steps", "步数"),
    ("done", "结束"),
)

FAILURE_FIELDS = (
    ("timestamp", "时间"),
    ("episode", "局"),
    ("failure_type", "失败模式"),
    ("evidence", "证据"),
    ("hypothesis", "假设"),
    ("next_step", "下一步"),
)

SPEED_LABELS = ("0.25x", "0.5x", "1x", "2x", "3x")


def read_jsonl(path: Path) -> list[dict]:
    if not path.exists():
        return []

    rows = []
    with path.open("r", encoding="utf-8") as handle:
        for line in handle:
            line = line.strip()
            if line:
                rows.append(json.loads(line))
    return rows


def read_csv(path: Path) -> list[dict]:
    if not path.exists():
        return []

    with path.open("r", encoding="utf-8", newline="") as handle:
        return list(csv.DictReader(handle))


def normalize_trials(trials: list[dict]) -> list[dict]:
    normalized = []
    for index, trial in enumerate(trials, start=1):
        row = dict(trial)
        row["round"] = int(row.get("trial_index") or index)
        row["suggestion"] = trial_suggestion(row)
        normalized.append(row)
    return normalized


def normalize_episodes(episodes: list[dict], trials: list[dict]) -> list[dict]:
    round_by_timestamp = {trial.get("timestamp"): trial["round"] for trial in trials}
    normalized = []
    for episode in episodes:
        row = dict(episode)
        row["round"] = int(row.get("trial_index") or round_by_timestamp.get(row.get("timestamp"), 0))
        normalized.append(row)
    return normalized


def normalize_failures(failures: list[dict], trials: list[dict]) -> list[dict]:
    round_by_timestamp = {trial.get("timestamp"): trial["round"] for trial in trials}
    normalized = []
    for failure in failures:
        row = dict(failure)
        row["round"] = int(row.get("trial_index") or round_by_timestamp.get(row.get("timestamp"), 0))
        normalized.append(row)
    return normalized


def trial_suggestion(trial: dict) -> str:
    try:
        score_mean = float(trial.get("score_mean", 0))
        score_min = float(trial.get("score_min", score_mean))
        score_max = float(trial.get("score_max", score_mean))
    except (TypeError, ValueError):
        return "先补证据，再改策略。"

    notes = str(trial.get("notes", "")).lower()
    if score_mean <= -20:
        return "仍全输；看 RGB 失败片段，并复核 ball_y / paddle_y 映射。"
    if score_max == score_min:
        return "分数无波动；每次只改一个启发式，再多 seed 测。"
    if "video" in notes or "gif" in notes:
        return "先用 RGB/OBS 标注失败模式，再改 heuristic。"
    return "和上一轮对比，每次只改一个 heuristic。"


def cell(value: object) -> str:
    return f"<td>{html.escape(str(value))}</td>"


def trial_summary_table(trials: list[dict]) -> str:
    rows = [
        "<tr>" + "".join(f"<th>{html.escape(label)}</th>" for _field, label in TRIAL_FIELDS) + "</tr>"
    ]
    for trial in reversed(trials):
        rows.append(
            "<tr>"
            + "".join(cell(trial.get(field, "")) for field, _label in TRIAL_FIELDS)
            + "</tr>"
        )
    return "<table>" + "\n".join(rows) + "</table>"


def round_sections(
    trials: list[dict],
    episodes: list[dict],
    failures: list[dict],
    report_path: Path,
) -> str:
    episodes_by_round: dict[int, list[dict]] = defaultdict(list)
    for episode in episodes:
        episodes_by_round[int(episode.get("round") or 0)].append(episode)

    failures_by_round: dict[int, list[dict]] = defaultdict(list)
    for failure in failures:
        failures_by_round[int(failure.get("round") or 0)].append(failure)

    sections = []
    for trial in reversed(trials):
        trial_round = int(trial["round"])
        trial_episodes = episodes_by_round.get(trial_round, [])
        trial_failures = failures_by_round.get(trial_round, [])
        sections.append(
            f"""
<section class="round">
  <div class="round-head">
    <div>
      <h3>Round {trial_round}: {html.escape(str(trial.get("policy_name", "")))}</h3>
      <div class="meta">{html.escape(str(trial.get("timestamp", "")))}</div>
    </div>
    <div class="score">均分 {html.escape(str(trial.get("score_mean", "")))}</div>
  </div>
  <dl class="round-facts">
    <dt>改动</dt><dd>{html.escape(str(trial.get("notes", "") or "未记录"))}</dd>
    <dt>样本</dt><dd>seed {html.escape(str(trial.get("seed", "")))} / {html.escape(str(trial.get("episodes", "")))} 局 / {html.escape(str(trial.get("env_steps", "")))} 步</dd>
    <dt>分数</dt><dd>{html.escape(str(trial.get("score_min", "")))} / {html.escape(str(trial.get("score_mean", "")))} / {html.escape(str(trial.get("score_max", "")))} 低/均/高</dd>
    <dt>建议</dt><dd>{html.escape(str(trial.get("suggestion", "")))}</dd>
  </dl>
  <h4>失败记录</h4>
  <div class="table-wrap compact">
    {failure_table(trial_failures, report_path)}
  </div>
  <h4>逐局记录</h4>
  <div class="table-wrap">
    {episode_table(trial_episodes, report_path)}
  </div>
</section>
"""
        )
    return "\n".join(sections)


def failure_table(failures: list[dict], report_path: Path) -> str:
    rows = [
        "<tr>"
        + "".join(f"<th>{html.escape(label)}</th>" for _field, label in FAILURE_FIELDS)
        + "<th>视频</th>"
        + "</tr>"
    ]
    if not failures:
        colspan = len(FAILURE_FIELDS) + 1
        rows.append(f'<tr><td colspan="{colspan}">这一轮还没有失败模式记录。</td></tr>')
        return "<table>" + "\n".join(rows) + "</table>"

    for failure in failures:
        cells = [cell(failure.get(field, "")) for field, _label in FAILURE_FIELDS]
        video_path = failure.get("video_path")
        if video_path:
            path = Path(str(video_path))
            href = html.escape(relative_href(report_path, path))
            cells.append(f'<td><a class="file-link" href="{href}">打开</a></td>')
        else:
            cells.append("<td></td>")
        rows.append("<tr>" + "".join(cells) + "</tr>")
    return "<table>" + "\n".join(rows) + "</table>"


def episode_table(episodes: list[dict], report_path: Path) -> str:
    rows = [
        "<tr>"
        + "".join(f"<th>{html.escape(label)}</th>" for _field, label in EPISODE_FIELDS)
        + "<th>RGB</th>"
        + "<th>RGB 速度</th>"
        + "<th>RGB 预览</th>"
        + "<th>OBS</th>"
        + "<th>OBS 速度</th>"
        + "<th>OBS 预览</th>"
        + "</tr>"
    ]
    if not episodes:
        colspan = len(EPISODE_FIELDS) + 6
        rows.append(f'<tr><td colspan="{colspan}">这一轮没有逐局视频记录。</td></tr>')
        return "<table>" + "\n".join(rows) + "</table>"

    for episode in episodes:
        cells = [cell(episode.get(field, "")) for field, _label in EPISODE_FIELDS]
        rgb_video_path = episode.get("rgb_video_path") or episode.get("video_path", "")
        obs_video_path = episode.get("obs_video_path", "")
        cells.extend(video_cells(report_path, rgb_video_path, "rgb episode gif"))
        cells.extend(video_cells(report_path, obs_video_path, "obs episode gif"))
        rows.append("<tr>" + "".join(cells) + "</tr>")
    return "<table>" + "\n".join(rows) + "</table>"


def relative_href(report_path: Path, target_path: Path) -> str:
    if not target_path.is_absolute():
        target_path = Path.cwd() / target_path
    try:
        relative = target_path.resolve().relative_to(report_path.parent.resolve()).as_posix()
        return f"/files/{relative}"
    except ValueError:
        return target_path.resolve().as_uri()


def speed_variant_path(source: Path, label: str) -> Path:
    safe_label = label.replace(".", "p")
    return source.with_name(f"{source.stem}_speed_{safe_label}{source.suffix}")


def speed_links(report_path: Path, video_path: Path) -> str:
    links = []
    for label in SPEED_LABELS:
        path = speed_variant_path(video_path, label)
        safe_video_path = html.escape(video_path.as_posix())
        safe_label = html.escape(label)
        if path.exists():
            href = html.escape(relative_href(report_path, path))
            links.append(f'<a class="speed-link ready" href="{href}">{safe_label}</a>')
        else:
            links.append(
                f'<button class="speed-link missing" '
                f'data-path="{safe_video_path}" '
                f'data-label="{safe_label}" '
                f'onclick="generateSpeed(this)">生成 {safe_label}</button>'
            )
    return "".join(links)


def video_cells(report_path: Path, video_path: str, alt: str) -> list[str]:
    if not video_path:
        return ["<td></td>", "<td></td>", "<td></td>"]

    path = Path(video_path)
    href = html.escape(relative_href(report_path, path))
    label = html.escape(path.name)
    speeds = speed_links(report_path, path)
    return [
        f'<td><a class="file-link" href="{href}" title="{label}">打开</a></td>',
        f"<td>{speeds}</td>",
        f'<td><img class="thumb" src="{href}" alt="{html.escape(alt)}"></td>',
    ]


def build_report(
    trials_path: Path,
    summary_path: Path,
    episodes_path: Path,
    output_path: Path,
    failures_path: Path = Path("failures.jsonl"),
) -> None:
    trials = read_jsonl(trials_path)
    if not trials:
        trials = read_csv(summary_path)
    trials = normalize_trials(trials)
    episodes = normalize_episodes(read_jsonl(episodes_path), trials)
    failures = normalize_failures(read_jsonl(failures_path), trials)

    output_path.parent.mkdir(parents=True, exist_ok=True)
    latest = trials[-1] if trials else {}
    document = f"""<!doctype html>
<html lang="zh-CN">
<head>
  <meta charset="utf-8">
  <meta name="viewport" content="width=device-width, initial-scale=1">
  <title>Pong RAM 策略实验</title>
  <style>
    :root {{
      color-scheme: light;
      font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif;
      background: #f6f7f9;
      color: #171a1f;
    }}
    body {{
      margin: 0;
      padding: 16px;
    }}
    h1, h2, h3, h4 {{
      margin: 0 0 14px;
      font-weight: 650;
    }}
    h1 {{
      font-size: 22px;
    }}
    h2 {{
      font-size: 16px;
      margin-top: 16px;
    }}
    h3 {{
      font-size: 15px;
      margin-bottom: 4px;
    }}
    h4 {{
      font-size: 13px;
      margin: 10px 0 6px;
    }}
    .meta {{
      color: #5b6270;
      margin-bottom: 10px;
      font-size: 12px;
    }}
    .table-wrap {{
      overflow-x: auto;
      background: #ffffff;
      border: 1px solid #d9dde5;
      border-radius: 6px;
    }}
    .table-wrap.compact {{
      margin-bottom: 2px;
    }}
    table {{
      width: 100%;
      border-collapse: collapse;
      font-size: 12px;
    }}
    th, td {{
      border-bottom: 1px solid #e5e8ef;
      padding: 4px 6px;
      text-align: left;
      vertical-align: top;
      white-space: nowrap;
    }}
    th {{
      background: #eef1f5;
      font-weight: 650;
    }}
    tr:last-child td {{
      border-bottom: 0;
    }}
    a {{
      color: #1459c9;
    }}
    button {{
      font: inherit;
    }}
    .file-link {{
      font-weight: 650;
    }}
    .speed-link {{
      display: inline-block;
      margin: 0 3px 3px 0;
      border-radius: 4px;
      border: 1px solid #ccd3df;
      padding: 2px 5px;
      line-height: 1.2;
      text-decoration: none;
      background: #ffffff;
    }}
    .speed-link.ready {{
      border-color: #9db3d8;
      background: #eef5ff;
    }}
    .speed-link.missing {{
      color: #5b6270;
      cursor: pointer;
    }}
    .speed-link:disabled {{
      cursor: wait;
      opacity: 0.65;
    }}
    .thumb {{
      width: 112px;
      height: 112px;
      image-rendering: pixelated;
      object-fit: contain;
      background: #101318;
    }}
    .round {{
      margin-top: 10px;
      background: #ffffff;
      border: 1px solid #d9dde5;
      border-radius: 6px;
      padding: 10px;
    }}
    .round-head {{
      display: flex;
      justify-content: space-between;
      gap: 16px;
      align-items: flex-start;
    }}
    .score {{
      font-weight: 650;
      background: #eef1f5;
      border-radius: 6px;
      padding: 4px 8px;
      white-space: nowrap;
    }}
    .round-facts {{
      display: grid;
      grid-template-columns: 64px minmax(0, 1fr) 64px minmax(0, 1fr);
      gap: 4px 10px;
      margin: 8px 0 10px;
      font-size: 12px;
    }}
    dt {{
      color: #5b6270;
    }}
    dd {{
      margin: 0;
    }}
  </style>
</head>
<body>
  <h1>Pong RAM 策略实验</h1>
  <div class="meta">
    轮次: {len(trials)} · 局数: {len(episodes)} · 失败记录: {len(failures)} · 最新轮: {html.escape(str(latest.get("round", "")))}
  </div>

  <h2>全部轮次</h2>
  <div class="table-wrap">
    {trial_summary_table(trials)}
  </div>

  <h2>轮次详情</h2>
  {round_sections(trials, episodes, failures, output_path)}
  <script>
    async function generateSpeed(button) {{
      const path = button.dataset.path;
      const label = button.dataset.label;
      button.disabled = true;
      button.textContent = "生成中 " + label;
      try {{
        const params = new URLSearchParams({{ path, label }});
        const response = await fetch("/api/speed?" + params.toString(), {{ method: "POST" }});
        if (!response.ok) {{
          throw new Error(await response.text());
        }}
        await response.json();
        window.location.reload();
      }} catch (error) {{
        button.disabled = false;
        button.textContent = "失败 " + label;
        window.alert("生成失败: " + error.message);
      }}
    }}
  </script>
</body>
</html>
"""
    output_path.write_text(document, encoding="utf-8")


def main() -> None:
    parser = argparse.ArgumentParser(description="Build report.html from local trial logs.")
    parser.add_argument("--trials-path", type=Path, default=Path("trials.jsonl"))
    parser.add_argument("--summary-path", type=Path, default=Path("summary.csv"))
    parser.add_argument("--episodes-path", type=Path, default=Path("episodes.jsonl"))
    parser.add_argument("--failures-path", type=Path, default=Path("failures.jsonl"))
    parser.add_argument("--output", type=Path, default=Path("report.html"))
    args = parser.parse_args()

    build_report(
        trials_path=args.trials_path,
        summary_path=args.summary_path,
        episodes_path=args.episodes_path,
        failures_path=args.failures_path,
        output_path=args.output,
    )
    print(f"report written: {args.output}")


if __name__ == "__main__":
    main()
