"""
Produce stratified train/val/test splits from the master manifest.

Stratification preserves the per-class ratio across splits — important when
the classes are imbalanced (and they always are, eventually).

Determinism: a fixed seed plus sha256 sorting means re-running on the same
manifest produces identical splits. This matters for reproducibility.

Usage
-----
    python scripts/dataset/split.py \
        --manifest data/manifest.csv \
        --out-dir data \
        --val 0.1 --test 0.1

Generator-aware Stage 3A split:

    python scripts/dataset/split.py \
        --manifest data/manifest.csv \
        --out-dir data \
        --stratify-by class-generator

Held-out generator evaluation split:

    python scripts/dataset/split.py \
        --manifest data/manifest.csv \
        --out-dir data \
        --stratify-by class-generator \
        --holdout-generator sdxl
"""
from __future__ import annotations

import argparse
import csv
import random
from collections import defaultdict
from pathlib import Path


def _generator_label(row: dict[str, str]) -> str:
    """Return the best available generator label for an AI row."""
    return row.get("generator") or row.get("source") or "unknown"


def _split_key(row: dict[str, str], stratify_by: str) -> str:
    """Return the grouping key used for stratified splitting."""
    cls = row["class"]
    if stratify_by == "class-generator" and cls == "ai_generated":
        return f"{cls}:{_generator_label(row)}"
    return cls


def _write_csv(path: Path, fieldnames: list[str], rows: list[dict[str, str]]) -> None:
    with path.open("w", newline="") as fh:
        writer = csv.DictWriter(fh, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--manifest", type=Path, required=True)
    parser.add_argument("--out-dir", type=Path, required=True)
    parser.add_argument("--val", type=float, default=0.1)
    parser.add_argument("--test", type=float, default=0.1)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument(
        "--stratify-by",
        choices=["class", "class-generator"],
        default="class",
        help=(
            "Stratification mode. 'class' preserves legacy behavior; "
            "'class-generator' additionally balances AI generators."
        ),
    )
    parser.add_argument(
        "--holdout-generator",
        action="append",
        default=[],
        help=(
            "AI generator to remove from train/val/test and write to "
            "heldout.csv. Can be passed multiple times."
        ),
    )
    args = parser.parse_args()

    if args.val + args.test >= 1.0:
        raise ValueError("val + test must be < 1.0")

    with args.manifest.open() as fh:
        reader = csv.DictReader(fh)
        rows = list(reader)
        fieldnames = reader.fieldnames or []

    holdout_generators = set(args.holdout_generator)
    split_candidates: list[dict[str, str]] = []
    holdout_rows: list[dict[str, str]] = []
    for r in rows:
        if r["class"] == "ai_generated" and _generator_label(r) in holdout_generators:
            holdout_rows.append(r)
        else:
            split_candidates.append(r)

    if holdout_generators and not holdout_rows:
        raise ValueError(
            "No rows matched --holdout-generator values: "
            f"{sorted(holdout_generators)}"
        )

    by_group: dict[str, list[dict[str, str]]] = defaultdict(list)
    for r in split_candidates:
        by_group[_split_key(r, args.stratify_by)].append(r)

    rng = random.Random(args.seed)
    train_rows, val_rows, test_rows = [], [], []
    for group, items in sorted(by_group.items()):
        # Sort by sha256 for determinism, then shuffle with seeded RNG.
        items.sort(key=lambda r: r["sha256"])
        rng.shuffle(items)
        n = len(items)
        n_test = int(round(n * args.test))
        n_val = int(round(n * args.val))
        test_rows.extend(items[:n_test])
        val_rows.extend(items[n_test : n_test + n_val])
        train_rows.extend(items[n_test + n_val :])
        print(
            f"  {group}: total={n}  "
            f"train={n - n_test - n_val}  val={n_val}  test={n_test}"
        )

    args.out_dir.mkdir(parents=True, exist_ok=True)
    for name, rs in [("train", train_rows), ("val", val_rows), ("test", test_rows)]:
        path = args.out_dir / f"{name}.csv"
        _write_csv(path, fieldnames, rs)
        print(f"  wrote {path} ({len(rs)} rows)")

    if holdout_generators:
        path = args.out_dir / "heldout.csv"
        holdout_rows.sort(key=lambda r: (r["sha256"], r["path"]))
        _write_csv(path, fieldnames, holdout_rows)
        print(
            f"  wrote {path} ({len(holdout_rows)} rows) "
            f"for held-out generators: {sorted(holdout_generators)}"
        )


if __name__ == "__main__":
    main()