#!/usr/bin/env python3
"""Split extracted gate rows into prereq_op_selector accepted / rejected / schema-drift buckets.

Reads an NDJSON file of extracted gate rows (from extract_gate_rows.py) and
produces three output files covering only prereq_op_selector rows:

  accepted.ndjson     — rows with no rejection flags
  rejected.ndjson     — rows with any rejection flag
  schema_drift.ndjson — subset of rejected rows where rejection is caused by
                        invalid_prerequisite_values (schema-drift label pollution)

Also prints a brief summary to stdout.
"""

import argparse
import json
from pathlib import Path

SCHEMA_DRIFT_FLAG_PREFIX = "invalid_prerequisite_values"


def is_schema_drift(row: dict) -> bool:
    return any(f.startswith(SCHEMA_DRIFT_FLAG_PREFIX) for f in row.get("rejection_flags", []))


def write_ndjson(path: Path, rows: list[dict]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as fh:
        for row in rows:
            fh.write(json.dumps(row, sort_keys=True) + "\n")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--input",
        default="semantic/extracted_gate_rows_slice.ndjson",
        help="NDJSON file produced by extract_gate_rows.py.",
    )
    parser.add_argument(
        "--output-dir",
        default="semantic/prereq_op_dataset",
        help="Directory where accepted/rejected/schema_drift files are written.",
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    repo_root = Path(__file__).resolve().parent.parent
    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()

    accepted: list[dict] = []
    rejected: list[dict] = []
    schema_drift: list[dict] = []
    skipped = 0

    with input_path.open("r", encoding="utf-8") as fh:
        for line in fh:
            line = line.strip()
            if not line:
                continue
            row = json.loads(line)
            if row.get("gate_id") != "prereq_op_selector":
                skipped += 1
                continue
            flags = row.get("rejection_flags", [])
            if not flags:
                accepted.append(row)
            else:
                rejected.append(row)
                if is_schema_drift(row):
                    schema_drift.append(row)

    write_ndjson(output_dir / "accepted.ndjson", accepted)
    write_ndjson(output_dir / "rejected.ndjson", rejected)
    write_ndjson(output_dir / "schema_drift.ndjson", schema_drift)

    # Flag breakdown across rejected rows
    flag_counts: dict[str, int] = {}
    for row in rejected:
        for f in row.get("rejection_flags", []):
            flag_counts[f] = flag_counts.get(f, 0) + 1

    # Label distribution for accepted rows
    label_counts: dict[str, int] = {
        "needs_validate_intake": 0,
        "needs_architect_review": 0,
        "needs_resolve_dependencies": 0,
    }
    for row in accepted:
        label = row.get("decision_label", {})
        for key in label_counts:
            if label.get(key):
                label_counts[key] += 1

    summary = {
        "input": str(input_path),
        "output_dir": str(output_dir),
        "rows_accepted": len(accepted),
        "rows_rejected": len(rejected),
        "rows_schema_drift": len(schema_drift),
        "rows_skipped_other_gates": skipped,
        "rejection_flag_breakdown": flag_counts,
        "accepted_label_distribution": label_counts,
    }
    print(json.dumps(summary, indent=2, sort_keys=True))
    return 0


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