#!/usr/bin/env python3
"""Generate Fabricate-format TSV training data for the two prereq_op binary specialists.

Source: semantic/prereq_op_dataset_1000/accepted.ndjson
        (extracted from real whetstone_DSL taskitem run artifacts)

Produces four files:
  prereq_resolve_train.tsv   / prereq_resolve_eval.tsv
  prereq_architect_train.tsv / prereq_architect_eval.tsv

Each file uses Fabricate TSV format: label<TAB>0<TAB>text
  label: 0 = no, 1 = yes
  0:     unused hop field (always 0)
  text:  composed from task context fields

Text composition (per row):
  title: <task title>
  reasons: <reason1>; <reason2>; ...
  accepts: <acceptance criterion text>; ...
  constraints: <constraint text>; ...

Usage:
  python3 tools/gen_prereq_op_rsa_data.py \\
      --input semantic/prereq_op_dataset_1000/accepted.ndjson \\
      --output-dir semantic/prereq_op_tsv \\
      --eval-frac 0.15 \\
      --seed 42
"""

import argparse
import json
import random
from pathlib import Path


REPO_ROOT = Path(__file__).resolve().parent.parent


def compose_text(ctx: dict) -> str:
    parts = []

    title = (ctx.get("title") or "").strip()
    if title:
        parts.append(f"title: {title}")

    reasons = ctx.get("reasons") or []
    if reasons:
        parts.append("reasons: " + "; ".join(str(r).strip() for r in reasons if r))

    reqs = ctx.get("normalized_requirements") or []
    accepts = [r["normalizedText"] for r in reqs if r.get("kind") == "acceptance" and r.get("normalizedText")]
    constraints = [r["normalizedText"] for r in reqs if r.get("kind") == "constraint" and r.get("normalizedText")]

    if accepts:
        parts.append("accepts: " + "; ".join(accepts))
    if constraints:
        parts.append("constraints: " + "; ".join(constraints))

    return " | ".join(parts)


def write_tsv(path: Path, rows: list[tuple[int, str]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as fh:
        for label, text in rows:
            fh.write(f"{label}\t0\t{text}\n")


def label_summary(rows: list[tuple[int, str]]) -> dict:
    pos = sum(1 for label, _ in rows if label == 1)
    neg = len(rows) - pos
    return {"total": len(rows), "pos": pos, "neg": neg, "pos_pct": round(pos / len(rows) * 100, 1)}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--input",
        default="semantic/prereq_op_dataset_1000/accepted.ndjson",
        help="Accepted prereq_op rows in NDJSON format.",
    )
    parser.add_argument(
        "--output-dir",
        default="semantic/prereq_op_tsv",
        help="Directory for output TSV files.",
    )
    parser.add_argument(
        "--eval-frac",
        type=float,
        default=0.15,
        help="Fraction of rows to use for eval (default: 0.15).",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=42,
        help="Random seed for train/eval split.",
    )
    parser.add_argument(
        "--max-rows",
        type=int,
        default=0,
        help="Optional cap on input rows (0 = no limit).",
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    input_path = Path(args.input)
    if not input_path.is_absolute():
        input_path = (REPO_ROOT / input_path).resolve()
    output_dir = Path(args.output_dir)
    if not output_dir.is_absolute():
        output_dir = (REPO_ROOT / output_dir).resolve()

    rows: list[dict] = []
    with input_path.open("r", encoding="utf-8") as fh:
        for line in fh:
            line = line.strip()
            if not line:
                continue
            rows.append(json.loads(line))
            if args.max_rows and len(rows) >= args.max_rows:
                break

    rng = random.Random(args.seed)
    rng.shuffle(rows)

    n_eval = max(1, int(len(rows) * args.eval_frac))
    eval_rows = rows[:n_eval]
    train_rows = rows[n_eval:]

    def to_binary(split: list[dict], key: str) -> list[tuple[int, str]]:
        result = []
        for row in split:
            label_val = row["decision_label"].get(key, False)
            text = compose_text(row["decision_context"])
            if not text:
                continue
            result.append((1 if label_val else 0, text))
        return result

    resolve_train = to_binary(train_rows, "needs_resolve_dependencies")
    resolve_eval = to_binary(eval_rows, "needs_resolve_dependencies")
    architect_train = to_binary(train_rows, "needs_architect_review")
    architect_eval = to_binary(eval_rows, "needs_architect_review")

    write_tsv(output_dir / "prereq_resolve_train.tsv", resolve_train)
    write_tsv(output_dir / "prereq_resolve_eval.tsv", resolve_eval)
    write_tsv(output_dir / "prereq_architect_train.tsv", architect_train)
    write_tsv(output_dir / "prereq_architect_eval.tsv", architect_eval)

    summary = {
        "input": str(input_path),
        "output_dir": str(output_dir),
        "total_rows": len(rows),
        "train_rows": len(train_rows),
        "eval_rows": len(eval_rows),
        "needs_resolve_dependencies": {
            "train": label_summary(resolve_train),
            "eval": label_summary(resolve_eval),
        },
        "needs_architect_review": {
            "train": label_summary(architect_train),
            "eval": label_summary(architect_eval),
        },
    }
    print(json.dumps(summary, indent=2, sort_keys=True))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())