#!/usr/bin/env python3
"""Builds history_index.json for the interactive contestant history explorer (history.html)."""

from datetime import datetime, timezone
import json
import os
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent


def parse_frame_ts(filename: str) -> float | None:
    stem = os.path.splitext(filename)[0]
    try:
        return float(stem)
    except ValueError:
        pass
    try:
        iso = stem.replace("Z", "+00:00")
        # Python <3.11 / nanoseconds handling: trim fractional seconds to 6 digits if needed
        if "." in iso:
            head, tail = iso.split(".", 1)
            tz_idx = -1
            for sep in ("+", "-"):
                idx = tail.find(sep)
                if idx != -1:
                    tz_idx = idx
                    break
            if tz_idx != -1:
                frac, tz_part = tail[:tz_idx], tail[tz_idx:]
            else:
                frac, tz_part = tail, "+00:00"
            frac = (frac + "000000")[:6]
            iso = f"{head}.{frac}{tz_part}"
        return datetime.fromisoformat(iso).timestamp()
    except Exception:
        return None


def load_report_incidents(report_path: Path) -> dict[str, list[dict]]:
    by_user: dict[str, list[dict]] = {}
    if not report_path.exists():
        return by_user
    html = report_path.read_text(encoding="utf-8", errors="replace")
    marker = "const incidents = "
    idx = html.find(marker)
    if idx == -1:
        return by_user
    eps, _ = json.JSONDecoder().raw_decode(html, idx + len(marker))
    for ep in eps:
        cid = ep.get("contestant_id", "")
        ts_str = ep.get("timestamp", "")
        ts_epoch = None
        if ts_str:
            try:
                iso = ts_str.replace("Z", "+00:00")
                if "." in iso:
                    head, tail = iso.split(".", 1)
                    tz_idx = max(tail.find("+"), tail.find("-"))
                    frac = (tail[:tz_idx] + "000000")[:6] if tz_idx != -1 else (tail + "000000")[:6]
                    tz_part = tail[tz_idx:] if tz_idx != -1 else "+00:00"
                    iso = f"{head}.{frac}{tz_part}"
                ts_epoch = datetime.fromisoformat(iso).timestamp()
            except Exception:
                ts_epoch = None
        by_user.setdefault(cid, []).append(
            {
                "id": ep.get("id"),
                "timestamp": ts_str,
                "ts": ts_epoch,
                "severity": ep.get("severity"),
                "category": ep.get("category"),
                "confidence_score": ep.get("confidence_score"),
                "summary": ep.get("summary"),
                "details": ep.get("details"),
                "evidence_frame": ep.get("evidence_frame"),
                "context_frames": ep.get("context_frames", []),
                "code_diff_snippet": ep.get("code_diff_snippet"),
            }
        )
    for lst in by_user.values():
        lst.sort(key=lambda x: x["ts"] or 0.0)
    return by_user


def load_cms_export(cms_dir: Path):
    contest_json = cms_dir / "contest.json"
    if not contest_json.exists():
        return {}, None, None, [], None
    cj = json.loads(contest_json.read_text(encoding="utf-8"))

    group_start = None
    group_stop = None
    contest_tz = None
    tasks = {}
    users = {}
    participations = {}
    sub_results = {}

    for k, obj in cj.items():
        if not isinstance(obj, dict):
            continue
        cls = obj.get("_class")
        if cls == "Contest":
            contest_tz = obj.get("timezone")
        elif cls == "Group":
            group_start = obj.get("start")
            group_stop = obj.get("stop")
        elif cls == "Task":
            tasks[k] = {
                "id": k,
                "name": obj.get("name", k),
                "title": obj.get("title", obj.get("name", k)),
                "num": obj.get("num", 0),
            }
        elif cls == "User":
            users[k] = {
                "username": obj.get("username", ""),
                "first_name": obj.get("first_name", ""),
                "last_name": obj.get("last_name", ""),
            }
        elif cls == "Participation":
            participations[k] = {
                "user_id": obj.get("user"),
                "delay_time": obj.get("delay_time") or 0.0,
                "extra_time": obj.get("extra_time") or 0.0,
                "starting_time": obj.get("starting_time"),
                "hidden": bool(obj.get("hidden")),
            }
        elif cls == "SubmissionResult":
            sub_results[k] = {
                "score": obj.get("score"),
                "compilation_outcome": obj.get("compilation_outcome"),
                "compilation_stderr": (obj.get("compilation_stderr") or "")[:600],
            }

    per_username = {}
    for pid, part in participations.items():
        u = users.get(str(part["user_id"]), {})
        uname = u.get("username")
        if not uname:
            continue
        u_start = part["starting_time"] if part["starting_time"] is not None else (
            (group_start + part["delay_time"]) if group_start else None
        )
        u_stop = (
            (group_stop + part["delay_time"] + part["extra_time"])
            if group_stop
            else None
        )
        per_username[uname] = {
            "full_name": f"{u.get('first_name', '')} {u.get('last_name', '')}".strip(),
            "contest_start": u_start,
            "contest_stop": u_stop,
            "submissions": [],
            "questions": [],
        }

    for k, obj in cj.items():
        if not isinstance(obj, dict):
            continue
        cls = obj.get("_class")
        if cls == "Submission":
            part = participations.get(str(obj.get("participation")), {})
            uname = users.get(str(part.get("user_id")), {}).get("username")
            if not uname or uname not in per_username:
                continue
            task_info = tasks.get(str(obj.get("task")), {})
            res_ids = obj.get("results") or []
            score = None
            comp_outcome = None
            comp_err = ""
            for rid in res_ids:
                r = sub_results.get(str(rid))
                if r:
                    score = r.get("score")
                    comp_outcome = r.get("compilation_outcome")
                    comp_err = r.get("compilation_stderr") or ""
                    break
            files_out = []
            for fname_tpl, fid in (obj.get("files") or {}).items():
                fobj = cj.get(str(fid), {})
                digest = fobj.get("digest")
                if digest:
                    fpath = cms_dir / "files" / digest
                    size_bytes = fpath.stat().st_size if fpath.exists() else None
                    files_out.append(
                        {
                            "filename": obj.get("comment") or fobj.get("filename") or fname_tpl,
                            "url": str(fpath.relative_to(ROOT)),
                            "size": size_bytes,
                        }
                    )
            per_username[uname]["submissions"].append(
                {
                    "id": k,
                    "ts": obj.get("timestamp"),
                    "task": task_info.get("name", str(obj.get("task"))),
                    "task_title": task_info.get("title", ""),
                    "language": obj.get("language", ""),
                    "score": score,
                    "compilation_outcome": comp_outcome,
                    "compilation_stderr": comp_err,
                    "files": files_out,
                }
            )
        elif cls == "Question":
            part = participations.get(str(obj.get("participation")), {})
            uname = users.get(str(part.get("user_id")), {}).get("username")
            if not uname or uname not in per_username:
                continue
            per_username[uname]["questions"].append(
                {
                    "id": k,
                    "ts": obj.get("question_timestamp"),
                    "subject": obj.get("subject", ""),
                    "text": obj.get("text", ""),
                    "reply_subject": obj.get("reply_subject"),
                    "reply_text": obj.get("reply_text"),
                    "reply_ts": obj.get("reply_timestamp"),
                }
            )

    for info in per_username.values():
        info["submissions"].sort(key=lambda s: s["ts"] or 0.0)
        info["questions"].sort(key=lambda q: q["ts"] or 0.0)

    sorted_tasks = [v for _, v in sorted(tasks.items(), key=lambda kv: kv[1]["num"])]
    return per_username, group_start, group_stop, sorted_tasks, contest_tz


def estimate_client_clock_skew(c_dir: Path) -> float:
    logs_path = c_dir / "logs.json"
    if not logs_path.exists():
        return 0.0
    diffs = []
    for line in logs_path.read_text(encoding="utf-8", errors="replace").splitlines():
        line = line.strip()
        if not line:
            continue
        try:
            obj = json.loads(line)
            c_ts = float(obj["client_ts"])
            s_ts = float(obj["server_ts"])
            diffs.append(c_ts - s_ts)
        except Exception:
            continue
    if not diffs:
        return 0.0
    diffs.sort()
    idx = min(len(diffs) * 95 // 100, len(diffs) - 1)
    skew = diffs[idx]
    return skew if abs(skew) >= 5.0 else 0.0


def collect_contestant_frames(c_dir: Path) -> tuple[list[str], list[list], float]:
    skew = estimate_client_clock_skew(c_dir)
    dir_map: dict[str, int] = {}
    dir_list: list[str] = []
    frames = []
    for root, _, files in os.walk(c_dir, followlinks=True):
        rel_dir = os.path.relpath(root, ROOT).replace("\\", "/")
        for fn in files:
            low = fn.lower()
            if not low.endswith((".jpg", ".jpeg", ".png", ".webp")):
                continue
            ts = parse_frame_ts(fn)
            if ts is None:
                continue
            if rel_dir not in dir_map:
                dir_map[rel_dir] = len(dir_list)
                dir_list.append(rel_dir)
            frames.append([dir_map[rel_dir], fn, round(ts - skew, 3)])
    frames.sort(key=lambda x: (x[2], x[1]))
    return dir_list, frames, round(skew, 2)


def build_dataset(
    dataset_id: str,
    title: str,
    screenshots_root: Path,
    report_path: Path,
    cms_dir: Path | None = None,
    allowed_ids: set[str] | None = None,
) -> dict:
    incidents_by_user = load_report_incidents(report_path)
    cms_users, group_start, group_stop, tasks, contest_tz = (
        load_cms_export(cms_dir) if cms_dir else ({}, None, None, [], None)
    )

    contestants = []
    if screenshots_root.exists():
        for entry in sorted(screenshots_root.iterdir(), key=lambda p: p.name):
            if not entry.is_dir():
                continue
            cid = entry.name
            if allowed_ids is not None and cid not in allowed_ids:
                continue
            dirs, frames, skew_sec = collect_contestant_frames(entry)
            cache_file = entry / "frame_analyses_cache.json"
            cache_url = (
                os.path.relpath(cache_file, ROOT).replace("\\", "/")
                if cache_file.exists()
                else None
            )
            cms_u = cms_users.get(cid, {})
            subs = cms_u.get("submissions", [])
            best_by_task: dict[str, float] = {}
            for s in subs:
                sc = s.get("score")
                if sc is not None:
                    tname = s["task"]
                    best_by_task[tname] = max(best_by_task.get(tname, 0.0), float(sc))
            total_score = round(sum(best_by_task.values()), 2) if subs else None

            contestants.append(
                {
                    "id": cid,
                    "full_name": cms_u.get("full_name", ""),
                    "contest_start": cms_u.get("contest_start", group_start),
                    "contest_stop": cms_u.get("contest_stop", group_stop),
                    "clock_skew_sec": skew_sec,
                    "cache_url": cache_url,
                    "frame_dirs": dirs,
                    "frames": frames,
                    "submissions": subs,
                    "questions": cms_u.get("questions", []),
                    "incidents": incidents_by_user.get(cid, []),
                    "total_score": total_score,
                    "task_scores": best_by_task,
                }
            )

    return {
        "id": dataset_id,
        "title": title,
        "timezone": contest_tz or "UTC",
        "contest_start": group_start,
        "contest_stop": group_stop,
        "tasks": tasks,
        "report_url": os.path.relpath(report_path, ROOT).replace("\\", "/")
        if report_path.exists()
        else None,
        "contestants": contestants,
    }


def main():
    datasets = [
        build_dataset(
            dataset_id="volterra1",
            title="Volterra 1 — IOI 2027 Selection (90 Contestants)",
            screenshots_root=ROOT / "volterra1-ioi2027" / "oii-check-data",
            report_path=ROOT / "report_volterra1.html",
            cms_dir=ROOT / "volterra1-ioi2027" / "cms-export",
        ),
        build_dataset(
            dataset_id="tokens_2_8",
            title="Synthetic Benchmarks (token2 .. token8)",
            screenshots_root=ROOT / "screen-recording" / "screenshots",
            report_path=ROOT / "report_tokens_2_8.html",
            cms_dir=None,
            allowed_ids={f"token{i}" for i in range(2, 9)},
        ),
    ]
    out_path = ROOT / "history_index.json"
    out_path.write_text(
        json.dumps({"generated_at": datetime.now(timezone.utc).isoformat(), "datasets": datasets}, separators=(",", ":")),
        encoding="utf-8",
    )
    size_mb = out_path.stat().st_size / (1024 * 1024)
    total_frames = sum(len(c["frames"]) for d in datasets for c in d["contestants"])
    print(f"Wrote {out_path} ({size_mb:.2f} MB, {total_frames} frames indexed across {sum(len(d['contestants']) for d in datasets)} contestants)")


if __name__ == "__main__":
    main()