"""Phase 2 CLI: runs annotate_sample over the corpus, writing one JSONL
record per sample. Resumable (spec §4: "resumable... per-sample artifact
files, skip if exists") via a line-count check against an existing output
file combined with --resume; CPU-parallelizable via --workers (spec: "target
16 workers") using a process pool, since tree-sitter parsing is CPU-bound and
releases the GIL across process boundaries trivially via multiprocessing.

Does not invoke Joern (see src/cpg/annotate.py's module docstring) — this is
the tree-sitter-backed first pass. A real Joern batch pass + calibration
subset (spec §4.1's 200-sample edge-recall comparison) is tracked as
follow-up, not run yet.
"""
from __future__ import annotations

import argparse
import json
import signal
from multiprocessing import Pool
from pathlib import Path

import pandas as pd

from src.cpg.annotate import annotate_sample, annotation_to_dict

ROW_TIMEOUT_S = 15


class _RowTimeout(Exception):
    pass


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


def _annotate_row(row: dict) -> dict:
    """Never lets one pathological sample kill — or hang — the whole pool
    (spec §1.6: "fail loudly... never silently drop >2% of samples without
    logging a warning with counts and reasons"). Two failure modes guarded:
    exceptions (caught directly) and hangs/pathological slowness (bounded by
    a SIGALRM timeout). The latter was found empirically, not hypothetically:
    the real corpus contains at least one PrimeVul row whose "function" is a
    482KB/23,938-line blob (clearly a misclassified data table, not a real
    C function) that made build_defuse_graph's O(n^2) edge-building step
    (src/cpg/defuse.py) take so long it looked indistinguishable from a
    hang. A size threshold would have been a narrower, more brittle fix —
    the timeout catches this case AND any future pathological input we
    haven't seen yet, without needing to characterize what "too large" means
    in advance."""
    signal.signal(signal.SIGALRM, _alarm_handler)
    signal.alarm(ROW_TIMEOUT_S)
    try:
        ann = annotate_sample(
            sample_id=row["sample_id"],
            func_vuln=row["func_vuln"],
            patch_lines_vuln=list(row["patch_lines_vuln"]),
            cwe=row["cwe"],
        )
        return annotation_to_dict(ann)
    except _RowTimeout:
        return {
            "sample_id": row["sample_id"],
            "error": f"timeout after {ROW_TIMEOUT_S}s (func_len={len(row['func_vuln'])} chars)",
            "l1_lines": [], "l2_lines": [], "l3_lines": [],
            "trace_gt": [], "trace_unverifiable": True, "trace_reason": "annotation_timeout",
            "annotation_backend": "tree_sitter",
        }
    except Exception as e:  # noqa: BLE001 - deliberately broad, see docstring
        return {
            "sample_id": row["sample_id"],
            "error": f"{type(e).__name__}: {e}",
            "l1_lines": [], "l2_lines": [], "l3_lines": [],
            "trace_gt": [], "trace_unverifiable": True, "trace_reason": "annotation_error",
            "annotation_backend": "tree_sitter",
        }
    finally:
        signal.alarm(0)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--corpus", default="data/processed/corpus.parquet")
    parser.add_argument("--out", default="data/processed/annotations.jsonl")
    parser.add_argument("--workers", type=int, default=16)
    parser.add_argument("--limit", type=int, default=None, help="for smoke testing on a subset")
    parser.add_argument("--resume", action="store_true", default=False)
    args = parser.parse_args()

    df = pd.read_parquet(args.corpus)
    if args.limit:
        df = df.head(args.limit)

    out_path = Path(args.out)
    already_done: set[str] = set()
    if args.resume and out_path.exists():
        with open(out_path) as f:
            for line in f:
                line = line.strip()
                if line:
                    already_done.add(json.loads(line)["sample_id"])
        print(f"resuming: {len(already_done)} samples already annotated")

    rows = [r for r in df.to_dict("records") if r["sample_id"] not in already_done]
    print(f"annotating {len(rows)} samples with {args.workers} workers")

    out_path.parent.mkdir(parents=True, exist_ok=True)
    n_errors = 0
    n_trace_unverifiable = 0
    mode = "a" if args.resume else "w"
    with open(out_path, mode) as f, Pool(args.workers) as pool:
        for i, result in enumerate(pool.imap_unordered(_annotate_row, rows, chunksize=10)):
            f.write(json.dumps(result) + "\n")
            if result.get("error"):
                n_errors += 1
            elif result["trace_unverifiable"]:
                n_trace_unverifiable += 1
            if (i + 1) % 1000 == 0:
                f.flush()
                print(f"{i + 1}/{len(rows)} annotated ({n_errors} errors so far)")

    error_rate = n_errors / len(rows) if rows else 0.0
    print(f"done. {len(rows)} annotated, {n_trace_unverifiable} trace_unverifiable, "
          f"{n_errors} errors ({error_rate:.1%})")
    if error_rate > 0.02:
        print(f"WARNING: error rate {error_rate:.1%} exceeds the 2% threshold spec §1.6 "
              f"flags as needing investigation, not silent acceptance.")


if __name__ == "__main__":
    main()