"""Compare recorded CUA trajectories with real model rollouts on the GUI surface.

    python scripts/compare_cua.py --cua artifacts/cua/cua.jsonl --trace artifacts/cua-compare/trace.jsonl

For each side it reports how the episode was driven: turns, calls per turn,
the gesture mix, how numbers got into the editor (typed into fields vs dragged),
scrolls, undos, thought length, and reward. The point is to see whether the
recorded behaviour looks like what a capable model actually does, per layout.
"""

from __future__ import annotations

import argparse
import json
import statistics
from collections import Counter
from pathlib import Path


def _profile(turns: list[dict]) -> dict:
    """`turns`: [{"text": str, "calls": [{"name", "args"}]}] -> behaviour profile."""
    calls = [c for t in turns for c in t["calls"]]
    names = Counter(c["name"] for c in calls)
    typed = [c for c in calls if c["name"] == "type"]
    numeric = sum(1 for c in typed if str(c["args"].get("text", "")).replace(".", "").replace("-", "").isdigit())
    keys = Counter(str(c["args"].get("keys")) for c in calls if c["name"] == "key")
    words = [len(t["text"].split()) for t in turns if t.get("text")]
    return {
        "turns": len(turns),
        "calls": len(calls),
        "calls_per_turn": round(len(calls) / max(1, len(turns)), 2),
        "mix": {k: round(v / max(1, len(calls)), 3) for k, v in names.most_common()},
        "typed_numbers": numeric,
        "typed_text": len(typed) - numeric,
        "drags": names.get("left_click_drag", 0),
        "scrolls": names.get("scroll", 0),
        "undo": keys.get("ctrl+z", 0) + keys.get("cmd+z", 0),
        "screenshots_requested": names.get("screenshot", 0),
        "thought_words_median": statistics.median(words) if words else 0,
        "turns_with_thought": round(len(words) / max(1, len(turns)), 2),
    }


def from_cua(traj: dict) -> dict:
    turns = []
    for m in traj["messages"]:
        if m["role"] == "assistant":
            turns.append({"text": m.get("content") or "", "calls": [
                {"name": c["function"]["name"], "args": json.loads(c["function"]["arguments"])}
                for c in m.get("tool_calls", [])]})
    meta = traj["meta"]
    return {"who": "recorder", "skin": meta["skin"], "layout": meta.get("layout"), "task_id": meta["task_id"],
            "reward": meta["reward"], **_profile(turns)}


def from_trace(rec: dict) -> dict:
    turns = [{"text": t.get("text", ""), "calls": t.get("calls", [])} for t in rec.get("turns", [])]
    return {"who": rec["model"], "skin": rec.get("skin"), "layout": "default", "task_id": rec.get("task_id"),
            "reward": rec.get("reward"), **_profile(turns)}


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("--cua", action="append", default=[], help="cua.jsonl, or a directory of trajectory .json files")
    ap.add_argument("--trace", action="append", default=[], help="rollout trace.jsonl")
    ap.add_argument("--out", default=None)
    args = ap.parse_args(argv)
    rows = []
    for src in args.cua:
        p = Path(src)
        items = ([json.loads(line) for line in p.read_text().splitlines() if line.strip()] if p.is_file()
                 else [json.loads(f.read_text()) for f in sorted(p.glob("*.json"))])
        rows += [from_cua(t) for t in items]
    for src in args.trace:
        rows += [from_trace(json.loads(line)) for line in Path(src).read_text().splitlines()
                 if line.strip() and json.loads(line).get("surface") == "gui"]
    cols = ["who", "skin", "reward", "turns", "calls", "calls_per_turn", "typed_numbers", "typed_text", "drags",
            "scrolls", "undo", "screenshots_requested", "thought_words_median"]
    print("| " + " | ".join(cols) + " | top gestures |")
    print("|" + "---|" * (len(cols) + 1))
    for r in sorted(rows, key=lambda r: (r["skin"] or "", r["who"])):
        top = ", ".join(f"{k} {v:.0%}" for k, v in list(r["mix"].items())[:4])
        vals = [f"{r[c]:.3f}" if isinstance(r[c], float) and c == "reward" else str(r[c]) for c in cols]
        print("| " + " | ".join(vals) + f" | {top} |")
    if args.out:
        Path(args.out).write_text(json.dumps(rows, indent=1))
    return 0


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