"""Convert the harness trajectory export (JSONL) into the compact JSON the replay Space loads.

Usage:
    python fetch_trees.py                                          # once: metadata + repository trees
    python build_data.py minicpm5_2b_rl_fileline_turn02_trajectories.jsonl

Input: one or more JSONL files, one trajectory per line, schema_version 1 of the
evaluation harness export: {trajectory_id, model, benchmark, seed, instance_id, repo,
language, evaluation, messages, final_prediction, gold, scores, outcome, usage}.
``messages`` is the OpenAI-style chat transcript (system, user, assistant with
tool_calls, tool). Unlike the earlier grep/colgrep export, this one carries the gold
spans, the harness's parsed prediction and its file- and line-level scores, so none of
them are recomputed here. Base commits and repository trees still come from
``raw/instances.json`` and ``raw/trees/`` (see fetch_trees.py).

Output (in ./data):
    index.json          one row per instance with summary metrics
    traj/<id>.json      full replay for one instance: {..., "run": {...}}

Only seed 1 is kept so that each instance has one trajectory.
"""

from __future__ import annotations

import json
import pathlib
import re
import sys

# ---------------------------------------------------------------- configuration
HERE = pathlib.Path(__file__).parent
OUT_DIR = HERE / "data"
RAW_DIR = HERE / "raw"
KEEP_SEEDS = {1}
MODEL_NAME = "colgrep-agent-2B"
MODEL_URL = "https://huggingface.co/lightonai/colgrep-agent-2B"

# Curated examples shown at the top of the picker. Order matters. The note says what to
# watch for; it is written from the recorded seed-1 trajectories, not from the paper.
FEATURED: list[tuple[str, str]] = [
    (
        "django__django-11099",
        "Three turns, exact lines (Python). One query, one read of validators.py, then two single-line ranges: "
        "lines 10 and 20, the two regex literals the fix changes.",
    ),
    (
        "pydata__xarray-4248",
        "File found, lines missed (Python). formatting.py is the first result of the first query and the agent reads it three times, "
        "but at lines 1-100, 386-391 and 495-515. It submits dataset_repr (495-515); the gold line is 263.",
    ),
    (
        "matplotlib__matplotlib-25332",
        "A miss (Python). The gold file cbook.py surfaces at turn 5 under a weakref query and is never opened. The agent repeats "
        "the _align_label_groups query at turns 6, 8 and 9, and the harness forces a finish at turn 10 on figure.py.",
    ),
    (
        "cli__cli-405",
        "Two gold files (Go). The first query, printPrPreview, lists pr.go and issue.go. The agent reads the two preview "
        "functions, confirms the RenderMarkdown call with a second query, and submits one line in each file, both inside the gold spans.",
    ),
    (
        "serde-rs__serde-2798",
        "Two gold files (Rust). One query lists ser.rs and check.rs with the right ranges; two reads and a finish at turn 4. "
        "The submitted ranges fall inside the gold spans but cover a fifth of their lines (line recall 0.22).",
    ),
    (
        "darkreader__darkreader-7241",
        "Exact line hit in 4 turns (TypeScript). The first query's top result is parse.ts:146, the gold line. "
        "The agent reads the file twice to confirm and submits that single line.",
    ),
    (
        "jqlang__jq-1793",
        "Exact line hit (C). The first query lists jv.c:1336-1352 first, but the agent opens builtin.c first, then jv.c, "
        "and submits line 1345, the gold line.",
    ),
    (
        "fmtlib__fmt-2394",
        "One gold file, five gold spans (C++). The query on the emphasis enum lands on color.h:186. The agent reads the enum and "
        "the escape-sequence builder and submits three ranges covering most of the gold lines (line F1 0.81).",
    ),
    (
        "anuraghazra__github-readme-stats-1378",
        "Two gold files (JavaScript). The first query lists both; the agent reads them whole, then narrows with a second query on "
        "wrapTextMultiline. It submits the right block of utils.js but the wrong line of repo-card.js.",
    ),
    (
        "google__gson-1555",
        "Two gold files, eight gold spans (Java). Both files appear in the first listing. Three whole-file reads and a second query "
        "give three submitted ranges; file F1 1.0, line F1 0.38.",
    ),
]
# ------------------------------------------------------------------------------


def parse_args(raw: str) -> dict:
    try:
        return json.loads(raw)
    except json.JSONDecodeError:
        return {"raw": raw}


def describe_call(tool_name: str, args: dict) -> dict:
    """Classify a tool call for display: search, read, or finish."""
    if tool_name == "finish":
        locs = args.get("locations") or []
        return {"kind": "finish", "locations": [str(l) for l in locs]}
    if tool_name == "colgrep":
        return {
            "kind": "search",
            "tool": "colgrep",
            "query": args.get("query"),
            "pattern": args.get("pattern"),
            "path": args.get("path"),
            "include": args.get("include"),
            "k": args.get("k"),
            "command": args.get("command"),
        }
    cmd = (args.get("command") or args.get("raw") or "").strip()
    first = cmd.split()[0] if cmd else ""
    if first in {"grep", "rg", "find", "ag", "colgrep"}:
        return {"kind": "search", "tool": first, "command": cmd}
    return {"kind": "read", "tool": "terminal", "command": cmd}


PATH_TOKEN = re.compile(
    r"(?<![\w/])\.?/?((?:[\w.-]+/)*[\w.-]+\.[A-Za-z0-9]{1,8})(?![\w/])"
)
# colgrep result header: "[3] path/to/file.py:120-141  (function name)"
RESULT_LINE = re.compile(r"^\[\d+\] (\S+?):(\d+)-(\d+)", re.M)
# cat -n output line: "   120\t..."
NUMBERED_LINE = re.compile(r"^\s*(\d+)\t", re.M)


def paths_in(text: str, tree_paths: set[str]) -> list[str]:
    """Repository files mentioned in a text, in order of first appearance."""
    seen: list[str] = []
    for m in PATH_TOKEN.finditer(text):
        p = m.group(1).removeprefix("./")
        if p in tree_paths and p not in seen:
            seen.append(p)
    return seen


def span_overlap(a: tuple[int, int], b: tuple[int, int]) -> bool:
    return a[0] <= b[1] and b[0] <= a[1]


def spans_in_search_output(text: str) -> list[tuple[str, int, int]]:
    return [
        (m.group(1), int(m.group(2)), int(m.group(3)))
        for m in RESULT_LINE.finditer(text)
    ]


def lines_in_read_output(text: str) -> tuple[int, int] | None:
    nums = [int(m.group(1)) for m in NUMBERED_LINE.finditer(text)]
    return (min(nums), max(nums)) if nums else None


def gold_span_hits(
    spans: list[tuple[str, int, int]], gold_spans: list[dict]
) -> list[str]:
    """Gold spans (as 'file:start-end') overlapped by any of the given spans."""
    hits = []
    for g in gold_spans:
        key = f"{g['file']}:{g['start_line']}-{g['end_line']}"
        for f, lo, hi in spans:
            if f == g["file"] and span_overlap(
                (lo, hi), (g["start_line"], g["end_line"])
            ):
                hits.append(key)
                break
    return hits


def load_tree(instance_id: str) -> list[str]:
    f = RAW_DIR / "trees" / f"{instance_id}.json"
    return json.load(open(f))["paths"] if f.exists() else []


def dir_skeleton(paths: list[str], shown: set[str]) -> dict:
    """Nested directory skeleton with recursive file counts.

    Every directory is kept (with its file count) so the viewer can show the shape
    of the repository; files are listed only when they are in ``shown`` (gold,
    searched, read, or submitted), which keeps the payload small.
    """
    root: dict = {"n": 0, "d": {}, "f": []}
    for p in paths:
        parts = p.split("/")
        node = root
        node["n"] += 1
        for part in parts[:-1]:
            node = node["d"].setdefault(part, {"n": 0, "d": {}, "f": []})
            node["n"] += 1
        if p in shown:
            node["f"].append(parts[-1])

    def prune(node: dict) -> dict:
        out = {"n": node["n"]}
        if node["d"]:
            out["d"] = {k: prune(v) for k, v in sorted(node["d"].items())}
        if node["f"]:
            out["f"] = sorted(node["f"])
        return out

    return prune(root)


def split_problem(user_message: str) -> str:
    """Problem statement without the harness prefix, the repository listing and the final instruction."""
    stmt = user_message
    for marker in (
        "\n\nYou are at the repository root.",
        "\n\nFind the code locations relevant to this problem.",
    ):
        if marker in stmt:
            stmt = stmt.split(marker)[0]
    return stmt.removeprefix("Problem statement:\n").strip()


def convert_run(doc: dict, tree_paths: set[str]) -> dict:
    gold_files: list[str] = list(doc["gold"]["files"])
    gold_spans: list[dict] = list(doc["gold"]["spans"])
    turns: list[dict] = []
    system_prompt = ""
    problem = ""
    found_files: set[str] = set()
    found_spans: set[str] = set()
    pending_notes: list[str] = []  # harness messages addressed to the next turn
    for m in doc["messages"]:
        if m["role"] == "system":
            system_prompt = m["content"]
        elif m["role"] == "user":
            if not problem:
                problem = m["content"]
            else:
                pending_notes.append((m.get("content") or "").strip())
        elif m["role"] == "assistant":
            calls = [
                describe_call(
                    tc["function"]["name"], parse_args(tc["function"]["arguments"])
                )
                for tc in m.get("tool_calls") or []
            ]
            read_paths: list[str] = []
            for c in calls:
                if c["kind"] == "read" and c.get("command"):
                    read_paths += [
                        p
                        for p in paths_in(c["command"], tree_paths)
                        if p not in read_paths
                    ]
            turns.append(
                {
                    "n": len(turns) + 1,
                    "notes": pending_notes,
                    "thought": (m.get("content") or "").strip(),
                    "calls": calls,
                    "observations": [],
                    "touched": {"searched": [], "read": read_paths},
                }
            )
            pending_notes = []
        elif m["role"] == "tool" and turns:
            text = m.get("content") or ""
            # Strip the harness turn-counter footer; the viewer numbers turns itself.
            text = re.sub(
                r"\n*\[This is the output of turn \d+ \(\d+ turn limit\)\.\]\s*$",
                "",
                text,
            )
            k = len(turns[-1]["observations"])
            call = turns[-1]["calls"][k] if k < len(turns[-1]["calls"]) else None
            file_hits = [g for g in gold_files if g in text]
            new_files = [g for g in file_hits if g not in found_files]
            found_files.update(file_hits)
            # Gold lines: a search result range or a numbered read range that overlaps a gold span.
            span_hits: list[str] = []
            if call and call["kind"] == "search":
                span_hits = gold_span_hits(spans_in_search_output(text), gold_spans)
            elif call and call["kind"] == "read":
                rng = lines_in_read_output(text)
                files = turns[-1]["touched"]["read"]
                if rng and len(files) == 1:
                    span_hits = gold_span_hits([(files[0], rng[0], rng[1])], gold_spans)
            new_spans = [s for s in span_hits if s not in found_spans]
            found_spans.update(span_hits)
            turns[-1]["observations"].append(
                {
                    "text": text,
                    "gold_hits": file_hits,
                    "new_gold": new_files,
                    "span_hits": span_hits,
                    "new_span": new_spans,
                }
            )
            # Files surfaced by this observation (search hits or listings). Only
            # observations of search calls count as "searched"; read output is the
            # content of a file already recorded under "read".
            if call and call["kind"] == "search":
                t = turns[-1]["touched"]
                t["searched"] += [
                    p
                    for p in paths_in(text, tree_paths)
                    if p not in t["searched"] and p not in t["read"]
                ]
    if pending_notes and turns:
        turns[-1]["notes"] = turns[-1]["notes"] + pending_notes

    pred = doc["final_prediction"]
    locations = [
        {"file": l["file"], "start": l["start_line"], "end": l["end_line"]}
        for l in pred.get("locations") or []
    ]
    pred_files: list[str] = []
    for l in locations:
        if l["file"] not in pred_files:
            pred_files.append(l["file"])
    finish_turn = next(
        (t["n"] for t in turns if any(c["kind"] == "finish" for c in t["calls"])), None
    )
    forced = finish_turn is not None and any(
        "final turn" in n.lower() for n in turns[finish_turn - 1]["notes"]
    )
    s = doc["scores"]
    touched_all = {
        p for t in turns for k in ("searched", "read") for p in t["touched"][k]
    }
    return {
        "search_tool": doc["evaluation"]["search_tool"],
        "model": doc["model"],
        "evaluation": doc["evaluation"],
        "seed": doc["seed"],
        "problem": split_problem(problem),
        "system_prompt": system_prompt,
        "gold": gold_files,
        "gold_spans": [
            {"file": g["file"], "start": g["start_line"], "end": g["end_line"]}
            for g in gold_spans
        ],
        "predicted": pred_files,
        "locations": locations,
        "metrics": {
            "f1": round(s["file_f1"], 4),
            "precision": round(s["file_precision"], 4),
            "recall": round(s["file_recall"], 4),
            "line_f1": None if s["line_f1"] is None else round(s["line_f1"], 4),
            "line_precision": None
            if s["line_precision"] is None
            else round(s["line_precision"], 4),
            "line_recall": None
            if s["line_recall"] is None
            else round(s["line_recall"], 4),
        },
        "n_turns": len(turns),
        "n_searches": sum(
            1 for t in turns for c in t["calls"] if c["kind"] == "search"
        ),
        "n_reads": sum(1 for t in turns for c in t["calls"] if c["kind"] == "read"),
        "n_files_touched": len(touched_all),
        "completion_tokens": doc["usage"].get("completion_tokens"),
        "end_reason": doc["outcome"]["end_reason"],
        "finish_turn": finish_turn,
        "forced_finish": forced,
        "turns": turns,
    }


def main(paths: list[str]) -> None:
    meta_by_id = json.load(open(RAW_DIR / "instances.json"))
    by_instance: dict[str, dict] = {}
    skipped = 0
    missing_meta: set[str] = set()
    for path in paths:
        with open(path) as fh:
            for line in fh:
                doc = json.loads(line)
                if doc["seed"] not in KEEP_SEEDS:
                    skipped += 1
                    continue
                iid = doc["instance_id"]
                if iid not in meta_by_id:
                    missing_meta.add(iid)
                    skipped += 1
                    continue
                md = meta_by_id[iid]
                if sorted(md["gold"]) != sorted(doc["gold"]["files"]):
                    print(
                        "gold files differ from raw/instances.json for",
                        iid,
                        file=sys.stderr,
                    )
                tree_paths = load_tree(iid)
                run = convert_run(doc, set(tree_paths))
                by_instance[iid] = {
                    "instance_id": iid,
                    "benchmark": doc["benchmark"],
                    "repo": md["repo"],
                    "github_repo": md["repo"],
                    "base_commit": md["base_commit"],
                    "language": md["language"],
                    "title": run["problem"].splitlines()[0].strip()
                    if run["problem"]
                    else iid,
                    "gold": run["gold"],
                    "gold_spans": run["gold_spans"],
                    "run": run,
                    "tree_paths": tree_paths,
                }

    (OUT_DIR / "traj").mkdir(parents=True, exist_ok=True)
    for old in (OUT_DIR / "traj").glob("*.json"):
        old.unlink()
    index = []
    for iid, inst in sorted(by_instance.items()):
        paths = inst.pop("tree_paths")
        r = inst["run"]
        if paths:
            shown = set(inst["gold"]) | set(r["predicted"])
            for t in r["turns"]:
                shown.update(t["touched"]["searched"])
                shown.update(t["touched"]["read"])
            inst["repo_tree"] = dir_skeleton(paths, shown)
            inst["n_files"] = len(paths)
        with open(OUT_DIR / "traj" / f"{iid}.json", "w") as fh:
            json.dump(inst, fh, ensure_ascii=False, separators=(",", ":"))
        row = {
            k: inst[k]
            for k in ("instance_id", "benchmark", "repo", "language", "title", "gold")
        }
        row["n_gold_spans"] = len(inst["gold_spans"])
        row["n_files"] = inst.get("n_files")
        row["f1"] = r["metrics"]["f1"]
        row["line_f1"] = r["metrics"]["line_f1"]
        row["n_turns"] = r["n_turns"]
        row["n_searches"] = r["n_searches"]
        row["n_reads"] = r["n_reads"]
        row["touched"] = r["n_files_touched"]
        row["forced_finish"] = r["forced_finish"]
        index.append(row)

    featured = [{"instance_id": i, "note": n} for i, n in FEATURED if i in by_instance]
    missing_featured = [i for i, _ in FEATURED if i not in by_instance]
    any_run = next(iter(by_instance.values()))["run"]
    meta = {
        "n_instances": len(index),
        "benchmarks": {
            b: sum(1 for i in index if i["benchmark"] == b)
            for b in sorted({i["benchmark"] for i in index})
        },
        "tool": any_run["search_tool"],
        "model_name": MODEL_NAME,
        "model_url": MODEL_URL,
        "model": any_run["model"],
        "evaluation": any_run["evaluation"],
        "seeds": sorted(KEEP_SEEDS),
        "with_tree": sum(1 for i in by_instance.values() if i.get("n_files")),
        "skipped_docs": skipped,
        "featured": featured,
    }
    with open(OUT_DIR / "index.json", "w") as fh:
        json.dump(
            {"meta": meta, "instances": index},
            fh,
            ensure_ascii=False,
            separators=(",", ":"),
        )
    print(json.dumps({k: v for k, v in meta.items() if k != "featured"}, indent=2))
    if missing_meta:
        print(
            "no metadata for",
            len(missing_meta),
            "instances:",
            sorted(missing_meta)[:10],
        )
    if missing_featured:
        print("featured instances not in the export:", missing_featured)
    for b in meta["benchmarks"]:
        rows = [i for i in index if i["benchmark"] == b]
        n = len(rows)
        print(
            f"{b:24s} mean file F1 {100 * sum(i['f1'] for i in rows) / n:5.1f}"
            f"  mean line F1 {100 * sum(i['line_f1'] or 0 for i in rows) / n:5.1f}"
            f"  forced finish {sum(i['forced_finish'] for i in rows):3d}  n={n}"
        )


if __name__ == "__main__":
    main(sys.argv[1:])