"""Build compressed, per-sample DFG caches for chronological-train R3 scoring."""
from __future__ import annotations

import argparse
import gzip
import json
import signal
from collections import Counter
from multiprocessing import Pool
from pathlib import Path

import pandas as pd

from src.rewards.checker import trace_validity
from src.rewards.cache import cache_path
from src.rewards.line_graph import build_line_graph

ROW_TIMEOUT_S = 15


class _RowTimeout(Exception):
    pass


def _alarm_handler(signum, frame):
    raise _RowTimeout()


def _build_one(task: tuple[dict, dict, str]) -> dict:
    row, annotation, root_name = task
    sample_id = row["sample_id"]
    signal.signal(signal.SIGALRM, _alarm_handler)
    signal.alarm(ROW_TIMEOUT_S)
    try:
        graph = build_line_graph(row["func_vuln"])
        sinks = set()
        for path in annotation["trace_gt"]:
            lines = [graph.node_lines[tuple(map(int, node.split("-", 1)))] for node in path]
            valid, _ = trace_validity(lines, graph, sink=lines[-1])
            if not valid:
                return {"sample_id": sample_id, "status": "trace_mismatch"}
            sinks.add(lines[-1])
        value = {"sample_id": sample_id, "annotation_backend": annotation["annotation_backend"],
                 "sink_lines": sorted(sinks), "graph": graph.to_dict()}
        path = cache_path(Path(root_name), sample_id)
        tmp = path.with_suffix(".tmp")
        with gzip.open(tmp, "wt") as f:
            json.dump(value, f)
        tmp.replace(path)
        return {"sample_id": sample_id, "status": "cached", "n_paths": len(annotation["trace_gt"])}
    except _RowTimeout:
        return {"sample_id": sample_id, "status": "timeout"}
    except Exception as exc:
        return {"sample_id": sample_id, "status": "error", "error": f"{type(exc).__name__}: {exc}"}
    finally:
        signal.alarm(0)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--corpus", default="data/processed/corpus.parquet")
    parser.add_argument("--split", default="data/splits/chronological/train.json")
    parser.add_argument("--annotations", default="data/processed/annotations.jsonl")
    parser.add_argument("--out-dir", default="data/cpg_cache/reward_graphs")
    parser.add_argument("--report", default="results/phase5/reward_graph_cache.json")
    parser.add_argument("--limit", type=int, default=None)
    parser.add_argument("--workers", type=int, default=4)
    parser.add_argument("--seed", type=int, default=42)
    args = parser.parse_args()

    train_pairs = set(json.loads(Path(args.split).read_text()))
    annotations = {a["sample_id"]: a for line in Path(args.annotations).read_text().splitlines()
                   if line.strip() and (a := json.loads(line)).get("trace_gt")}
    df = pd.read_parquet(args.corpus, columns=["sample_id", "pair_id", "target", "func_vuln"])
    rows = [row for row in df.to_dict("records") if row["pair_id"] in train_pairs
            and int(row["target"]) == 1 and row["sample_id"] in annotations]
    if args.limit is not None:
        rows = rows[:args.limit]
    root = Path(args.out_dir)
    root.mkdir(parents=True, exist_ok=True)
    existing = sum(cache_path(root, row["sample_id"]).exists() for row in rows)
    tasks = [(row, annotations[row["sample_id"]], str(root)) for row in rows
             if not cache_path(root, row["sample_id"]).exists()]
    print(f"trace-covered rows for split {args.split}={len(rows)}, existing={existing}, pending={len(tasks)}", flush=True)

    counts: Counter = Counter({"already_cached": existing})
    failures = []
    with Pool(max(1, args.workers)) as pool:
        for i, result in enumerate(pool.imap_unordered(_build_one, tasks, chunksize=1), 1):
            counts[result["status"]] += 1
            if result["status"] not in ("cached", "already_cached"):
                failures.append(result)
            if i % 500 == 0:
                print(f"processed {i}/{len(tasks)}; cached={counts['cached']}; errors={len(failures)}", flush=True)

    report = {"seed": args.seed, "split": args.split, "n_trace_covered_rows": len(rows),
              "counts": dict(counts), "n_cached_total": counts["already_cached"] + counts["cached"],
              "error_rate": len(failures) / len(rows) if rows else 0.0,
              "failures": failures[:50], "cache_dir": str(root)}
    report_path = Path(args.report)
    report_path.parent.mkdir(parents=True, exist_ok=True)
    report_path.write_text(json.dumps(report, indent=2) + "\n")
    print(json.dumps({k: v for k, v in report.items() if k != "failures"}, indent=2))
    if report["error_rate"] > 0.02:
        raise ValueError("reward graph cache error rate exceeds 2%")


if __name__ == "__main__":
    main()