#!/usr/bin/env python3
"""
gen_prereq_op_data.py

Mines whetstone_DSL sprint plan files to produce labeled training data for
the prereq_op_selector specialist.

Prerequisite op classes:
  standard       — [validate-intake, resolve-dependencies]
                   Default implementation task; no special gates needed.
  needs_review   — [validate-intake, architect-review]
                   Introduces new architecture, API surface, subsystem, or
                   cross-component interface. Requires architect sign-off.
  needs_approval — [validate-intake, manual-approval]
                   Touches release, security, compliance, or production.
                   Requires manual gate regardless of architecture complexity.
  full_gates     — [validate-intake, architect-review, manual-approval]
                   Both architectural and production-sensitive risk present.

Output: TSV with columns: label<TAB>hops<TAB>text
  label: integer index 0–3
         0=standard  1=needs_review  2=needs_approval  3=full_gates
  hops:  always 0 (unused by fabricate binary; required by format)
  text:  the step description text

Usage:
    python3 specialists/scripts/gen_prereq_op_data.py \\
        --sprint-dir /path/to/whetstone_DSL \\
        --train-out specialists/data/generated/prereq_op_train.tsv \\
        --eval-out  specialists/data/generated/prereq_op_eval.tsv
"""
import re
import sys
import random
import argparse
from pathlib import Path
from collections import Counter

# ---------------------------------------------------------------------------
# Classification heuristics
# ---------------------------------------------------------------------------

# Signals that an architect needs to review this (new API, structure, design)
REVIEW_TOKENS = {
    "introduce", "phase", "framework", "architecture", "design", "refactor",
    "planning", "taxonomy", "grammar", "interface", "api surface", "subsystem",
    "extension point", "abstraction", "decompos", "scaffold", "polyglot",
    "cross-language", "cross-project", "multi-language", "orchestrat",
    "migration plan", "phase plan", "protocol design", "new layer",
    "foundation", "bootstrap", "initial", "phase 1", "phase 2", "phase 3",
    "sprint plan", "roadmap",
}

# Signals that a manual approval gate is required (release, security, production)
APPROVAL_TOKENS = {
    "release", "deploy", "publish", "migrate", "security", "compliance",
    "upgrade", "breaking change", "production", "rollout", "finalize",
    "sign off", "certification", "approval", "cutover", "go-live",
    "data migration", "schema migration", "breaking", "deprecat",
}

LABEL_ORDER = ["standard", "needs_review", "needs_approval", "full_gates"]
LABEL_INDEX = {lbl: i for i, lbl in enumerate(LABEL_ORDER)}


def classify(text: str) -> str:
    t = text.lower()
    has_review   = any(tok in t for tok in REVIEW_TOKENS)
    has_approval = any(tok in t for tok in APPROVAL_TOKENS)

    if has_review and has_approval:
        return "full_gates"
    if has_review:
        return "needs_review"
    if has_approval:
        return "needs_approval"
    return "standard"


# ---------------------------------------------------------------------------
# Step extraction (same regex as gen_verification_type_data.py)
# ---------------------------------------------------------------------------

STEP_RE = re.compile(
    r"###\s+Step\s+\d+[a-z]?:\s+(.+?)(?:\s*\(\d+\s+tests?\))?$",
    re.IGNORECASE,
)


def extract_steps(sprint_path: Path) -> list[tuple[str, str]]:
    results = []
    for line in sprint_path.read_text(errors="replace").splitlines():
        m = STEP_RE.match(line.strip())
        if not m:
            continue
        desc = m.group(1).strip()
        if not desc or len(desc) < 8:
            continue
        label = classify(desc)
        results.append((label, desc))
    return results


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--sprint-dir", default=".",
                    help="Root of whetstone_DSL (contains sprint*_plan.md files)")
    ap.add_argument("--train-out", default="specialists/data/generated/prereq_op_train.tsv")
    ap.add_argument("--eval-out",  default="specialists/data/generated/prereq_op_eval.tsv")
    ap.add_argument("--eval-frac", type=float, default=0.15)
    ap.add_argument("--seed", type=int, default=42)
    args = ap.parse_args()

    sprint_dir = Path(args.sprint_dir)
    plans = sorted(sprint_dir.glob("sprint*_plan.md"))
    if not plans:
        print(f"No sprint plan files found in {sprint_dir}", file=sys.stderr)
        sys.exit(1)
    print(f"Found {len(plans)} sprint plan files")

    all_examples: list[tuple[str, str]] = []
    for p in plans:
        all_examples.extend(extract_steps(p))

    counts = Counter(lbl for lbl, _ in all_examples)
    print("Raw class distribution:", dict(counts))

    # Balance: cap at 3x the smallest class, floor at 40
    min_count = min(counts.values())
    target = min(min_count * 4, sorted(counts.values())[1] if len(counts) > 1 else min_count)
    target = max(target, 40)
    print(f"Balancing to {target} examples per class")

    rng = random.Random(args.seed)
    by_class: dict[str, list[str]] = {}
    for lbl, desc in all_examples:
        by_class.setdefault(lbl, []).append(desc)

    balanced: list[tuple[str, str]] = []
    for lbl in LABEL_ORDER:
        descs = by_class.get(lbl, [])
        if not descs:
            print(f"WARNING: no examples for class '{lbl}'", file=sys.stderr)
            continue
        rng.shuffle(descs)
        chosen = descs[:target]
        while len(chosen) < target:
            chosen += descs
        balanced.extend((lbl, d) for d in chosen[:target])

    rng.shuffle(balanced)
    print(f"Balanced total: {len(balanced)} examples")

    n_eval = max(1, int(len(balanced) * args.eval_frac))
    eval_set  = balanced[:n_eval]
    train_set = balanced[n_eval:]

    Path(args.train_out).parent.mkdir(parents=True, exist_ok=True)
    Path(args.eval_out).parent.mkdir(parents=True, exist_ok=True)

    def write_tsv(path, rows):
        with open(path, "w") as f:
            for lbl, desc in rows:
                idx = LABEL_INDEX[lbl]
                f.write(f"{idx}\t0\t{desc}\n")
        print(f"Wrote {len(rows)} rows → {path}")

    write_tsv(args.train_out, train_set)
    write_tsv(args.eval_out,  eval_set)

    for split_name, split in [("train", train_set), ("eval", eval_set)]:
        c = Counter(lbl for lbl, _ in split)
        print(f"  {split_name}: {dict(sorted(c.items()))}")


if __name__ == "__main__":
    main()