# load the hand-crafted seed jsonl, split into train/val/test, persist to
# data/pairs/. v1 does not download anything; the seed was produced by an
# off-line workflow under data/seed/synthetic_pairs.jsonl.
#
# outputs:
#   data/pairs/train.json   list of {"raw": str, "clean": str, "category": str}
#   data/pairs/val.json
#   data/pairs/test.json
#   data/pairs/meta.json    counts, seed, source path

import json
import random
from pathlib import Path
from typing import Optional

from cleanup.config import DataConfig


def _load_seed(path: Path) -> list[dict]:
    rows: list[dict] = []
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            obj = json.loads(line)
            if "raw" not in obj or "clean" not in obj:
                continue
            rows.append(
                {
                    "raw": obj["raw"],
                    "clean": obj["clean"],
                    "category": obj.get("category", "uncategorized"),
                }
            )
    return rows


def _split(rows: list[dict], splits, rng: random.Random) -> tuple[list, list, list]:
    # stratified by category so each split sees a balanced category mix.
    # this matters because counts per category are different (~70 to ~80
    # each) and we do not want a small category absent from val or test.
    by_cat: dict[str, list[dict]] = {}
    for r in rows:
        by_cat.setdefault(r["category"], []).append(r)
    train: list[dict] = []
    val: list[dict] = []
    test: list[dict] = []
    for cat, cat_rows in by_cat.items():
        rng.shuffle(cat_rows)
        n = len(cat_rows)
        n_val = max(1, int(round(n * splits.val)))
        n_test = max(1, int(round(n * splits.test)))
        n_train = n - n_val - n_test
        train += cat_rows[:n_train]
        val += cat_rows[n_train : n_train + n_val]
        test += cat_rows[n_train + n_val :]
    rng.shuffle(train)
    rng.shuffle(val)
    rng.shuffle(test)
    return train, val, test


def build_and_save(cfg: DataConfig, out_dir: Path, smoke: bool = False) -> dict:
    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    seed_path = Path(cfg.seed_path)
    if not seed_path.exists():
        raise FileNotFoundError(
            f"seed not found at {seed_path}. generate it via the synthetic-data "
            "workflow under handoffs/, or run scripts/01_download.py from a "
            "checkout that has data/seed/ populated."
        )

    rows = _load_seed(seed_path)
    if smoke:
        rows = rows[:200]
    print(f"[download] loaded {len(rows)} seed pairs from {seed_path}")

    rng = random.Random(cfg.random_seed)
    train, val, test = _split(rows, cfg.splits, rng)

    def _write(name: str, batch: list[dict]) -> int:
        path = out_dir / f"{name}.json"
        path.write_text(json.dumps(batch, ensure_ascii=False, indent=None))
        return len(batch)

    counts = {
        "train": _write("train", train),
        "val": _write("val", val),
        "test": _write("test", test),
    }

    meta = {
        "seed_path": str(seed_path),
        "random_seed": cfg.random_seed,
        "splits": cfg.splits.__dict__,
        "counts": counts,
        "smoke": smoke,
    }
    (out_dir / "meta.json").write_text(json.dumps(meta, indent=2))
    print(f"[download] wrote {counts} to {out_dir}")
    return meta


def load_pairs(data_dir, split: str, max_rows: Optional[int] = None) -> list[dict]:
    path = Path(data_dir) / f"{split}.json"
    rows = json.loads(path.read_text())
    if max_rows is not None and max_rows < len(rows):
        rows = rows[:max_rows]
    return rows