"""Convert CUAD (extractive-QA format) into multi-label clause-classification
training data.

CUAD ships as SQuAD-style question/answer-span data: per contract, ~41
questions, each answered with text spans (or none). To *train a classifier*
we need (clause_text -> set of category labels). This script:

  1. segments each CUAD contract into clauses (same heading logic as the app),
  2. labels a clause with category X if it overlaps a CUAD answer span for X's
     question (coverage >= --overlap),
  3. emits negatives (clauses with no label) so the model learns "none too",
  4. splits BY CONTRACT (never by clause) so the eval can't leak.

    python -m scripts.prepare_cuad                      (from backend/, needs data/cuad)
    -> data/cuad/train.jsonl, val.jsonl, labels.json    ({"text":..., "labels":[...]})

Only categories with real CUAD ground truth are trained (see CUAD_QUESTION_MAP);
the rest keep the rules/zero-shot path.
"""
from __future__ import annotations

import json
import pathlib
import sys

sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[1]))

from app.schema import Clause, SourceSpan  # noqa: E402
from app.segmentation import CLAUSE_START  # noqa: E402

ROOT = pathlib.Path(__file__).resolve().parents[2]
CUAD_JSON = ROOT / "data" / "cuad" / "CUAD_v1.json"
OUT_DIR = ROOT / "data" / "cuad"

# our category key -> substring identifying the CUAD question for it.
# (Only categories CUAD actually labels; indemnification/sla/payment etc. have
# no CUAD ground truth, so they stay on the rules/zero-shot path.)
CUAD_QUESTION_MAP = {
    "termination": "Termination For Convenience",
    "auto_renewal": "Renewal Term",
    "notice": "Notice Period To Terminate Renewal",
    "liability_cap": "Cap On Liability",
    "governing_law": "Governing Law",
    "exclusivity": "Exclusivity",
    "insurance": "Insurance",
    "audit_rights": "Audit Rights",
    "assignment": "Anti-Assignment",
    "ip_ownership": "Ip Ownership Assignment",
    "warranty": "Warranty Duration",
    "term": "Expiration Date",
}
LABELS = sorted(CUAD_QUESTION_MAP)


def segment_plain_text(text: str) -> list[Clause]:
    lines = text.split("\n")
    offsets, pos = [], 0
    for ln in lines:
        offsets.append(pos)
        pos += len(ln) + 1
    starts = [i for i, ln in enumerate(lines) if CLAUSE_START.match(ln)] or [0]
    out = []
    for n, li in enumerate(starts):
        s = offsets[li]
        e = offsets[starts[n + 1]] if n + 1 < len(starts) else len(text)
        out.append(Clause(id=f"c{n}", text=text[s:e],
                           span=SourceSpan(start_char=s, end_char=e, page=0)))
    return out


def coverage(gold: tuple[int, int], clause: tuple[int, int]) -> float:
    inter = max(0, min(gold[1], clause[1]) - max(gold[0], clause[0]))
    return inter / (gold[1] - gold[0]) if gold[1] > gold[0] else 0.0


def build(limit: int, overlap: float, val_frac: float):
    data = json.loads(CUAD_JSON.read_text())["data"][:limit]
    rows: list[tuple[int, dict]] = []  # (doc_index, example)

    for di, doc in enumerate(data):
        for para in doc["paragraphs"]:
            text = para["context"]
            clauses = segment_plain_text(text)
            gold: dict[str, list[tuple[int, int]]] = {k: [] for k in LABELS}
            for qa in para["qas"]:
                for key, q_sub in CUAD_QUESTION_MAP.items():
                    if q_sub in qa["question"]:
                        for ans in qa["answers"]:
                            s = ans["answer_start"]
                            gold[key].append((s, s + len(ans["text"])))
            for c in clauses:
                if len(c.text.strip()) < 25:
                    continue
                cspan = (c.span.start_char, c.span.end_char)
                labels = [k for k in LABELS
                          if any(coverage(g, cspan) >= overlap for g in gold[k])]
                rows.append((di, {"text": c.text.strip()[:4000], "labels": labels}))

    # contract-level split (deterministic — by doc index, no RNG)
    n_docs = len(data)
    val_docs = {i for i in range(n_docs) if i % max(1, round(1 / val_frac)) == 0}
    train = [ex for di, ex in rows if di not in val_docs]
    val = [ex for di, ex in rows if di in val_docs]
    return train, val


def main() -> None:
    if not CUAD_JSON.exists():
        sys.exit(f"CUAD not found at {CUAD_JSON}. Run scripts/download_cuad.py first.")
    limit, overlap, val_frac = 999, 0.3, 0.15
    for a in sys.argv:
        if a.startswith("--limit="):
            limit = int(a.split("=", 1)[1])
        if a.startswith("--overlap="):
            overlap = float(a.split("=", 1)[1])

    train, val = build(limit, overlap, val_frac)
    OUT_DIR.mkdir(parents=True, exist_ok=True)
    (OUT_DIR / "train.jsonl").write_text("".join(json.dumps(r) + "\n" for r in train))
    (OUT_DIR / "val.jsonl").write_text("".join(json.dumps(r) + "\n" for r in val))
    (OUT_DIR / "labels.json").write_text(json.dumps(LABELS, indent=2))

    pos = sum(1 for r in train if r["labels"])
    print(f"wrote {len(train)} train / {len(val)} val clauses "
          f"({pos} train clauses carry >=1 label) over {len(LABELS)} categories")
    print(f"  -> {OUT_DIR}/train.jsonl, val.jsonl, labels.json")


if __name__ == "__main__":
    main()