"""Phase 1 §3.3: builds the three split regimes (random, chronological,
repo-disjoint) plus the OOD manifest placeholder, at the pair level so the
pair-integrity constraint (§3.3.4: a pair_id never crosses partitions) holds
by construction.

Each split is written to data/splits/<regime>/{train,valid,test}.json as a
list of pair_id (not sample_id — both halves of a pair always travel
together, so pair_id is the right unit of assignment and the smaller file).
"""
from __future__ import annotations

import json
import random
from collections import defaultdict
from pathlib import Path

import pandas as pd

SPLITS_DIR = Path("data/splits")


def _pair_table(df: pd.DataFrame) -> pd.DataFrame:
    """One row per pair_id, carrying the fields splits are built from."""
    vuln_rows = df[df["target"] == 1].drop_duplicates("pair_id")
    return vuln_rows[["pair_id", "project", "source", "commit_date"]].reset_index(drop=True)


def build_random_split(
    pairs: pd.DataFrame, seed: int = 42, ratios: tuple[float, float, float] = (0.8, 0.1, 0.1)
) -> dict[str, list[str]]:
    """80/10/10 by pair, stratified by source so each split keeps a
    comparable primevul/secvuleval mix (spec §3.3.1)."""
    rng = random.Random(seed)
    splits: dict[str, list[str]] = {"train": [], "valid": [], "test": []}
    for _source, group in pairs.groupby("source"):
        ids = group["pair_id"].tolist()
        rng.shuffle(ids)
        n = len(ids)
        n_train = int(n * ratios[0])
        n_valid = int(n * ratios[1])
        splits["train"].extend(ids[:n_train])
        splits["valid"].extend(ids[n_train : n_train + n_valid])
        splits["test"].extend(ids[n_train + n_valid :])
    return splits


def build_chronological_split(
    pairs: pd.DataFrame,
    train_before: str,
    valid_before: str,
) -> tuple[dict[str, list[str]], dict]:
    """Train on commit_date < train_before, valid in [train_before, valid_before),
    test >= valid_before (spec §3.3.2). Only pairs with a known commit_date are
    eligible; pairs without one are returned separately as `undated` so callers
    can decide whether to fold them into train or exclude them (we exclude by
    default — undated data must not leak into a label that claims to be
    time-ordered)."""
    dated = pairs[pairs["commit_date"].notna()].copy()
    # format="mixed": commit_date mixes two real formats across sources —
    # PrimeVul/Big-Vul gives plain "YYYY-MM-DD", SecVulEval/GitHub gives full
    # ISO-8601 with time+"Z". Without format="mixed", pandas infers ONE format
    # from the data and silently coerces every row in the other format to
    # NaT — found empirically: this was dropping 1,891 valid PrimeVul-dated
    # pairs (all plain-date strings) from the chronological split entirely,
    # not just mis-reporting a coverage stat. See DECISIONS.md.
    dated["commit_date"] = pd.to_datetime(dated["commit_date"], utc=True, format="mixed", errors="coerce")
    dated = dated[dated["commit_date"].notna()]

    train_cut = pd.Timestamp(train_before, tz="UTC")
    valid_cut = pd.Timestamp(valid_before, tz="UTC")

    train_ids = dated[dated["commit_date"] < train_cut]["pair_id"].tolist()
    valid_ids = dated[(dated["commit_date"] >= train_cut) & (dated["commit_date"] < valid_cut)]["pair_id"].tolist()
    test_ids = dated[dated["commit_date"] >= valid_cut]["pair_id"].tolist()

    stats = {
        "n_total_pairs": len(pairs),
        "n_dated_pairs": len(dated),
        "n_undated_pairs_excluded": len(pairs) - len(dated),
        "n_train": len(train_ids),
        "n_valid": len(valid_ids),
        "n_test": len(test_ids),
        "train_before": train_before,
        "valid_before": valid_before,
    }
    return {"train": train_ids, "valid": valid_ids, "test": test_ids}, stats


def build_repo_disjoint_folds(
    pairs: pd.DataFrame, n_folds: int = 5, seed: int = 42
) -> list[dict[str, list[str]]]:
    """Assigns every pair of a project to exactly one fold (greedy balanced by
    vulnerable-sample count per project; `pairs` here is already the
    one-row-per-pair, target=1 table, so "vulnerable-sample count" == pair
    count for that project). Returns `n_folds` {train, valid, test} splits,
    each using one fold as test, the next as valid, the rest as train
    (spec §3.3.3: "5 folds stored for cross-validation")."""
    by_project: dict[str, list[str]] = defaultdict(list)
    for row in pairs.itertuples():
        by_project[row.project].append(row.pair_id)

    rng = random.Random(seed)
    projects = list(by_project.keys())
    rng.shuffle(projects)
    projects.sort(key=lambda p: -len(by_project[p]))  # largest first, greedy balance

    fold_sizes = [0] * n_folds
    fold_projects: list[list[str]] = [[] for _ in range(n_folds)]
    for project in projects:
        i = min(range(n_folds), key=lambda k: fold_sizes[k])
        fold_projects[i].append(project)
        fold_sizes[i] += len(by_project[project])

    fold_pair_ids = [
        [pid for project in fp for pid in by_project[project]] for fp in fold_projects
    ]

    folds = []
    for k in range(n_folds):
        test_ids = fold_pair_ids[k]
        valid_ids = fold_pair_ids[(k + 1) % n_folds]
        train_ids = [
            pid for j in range(n_folds) if j not in (k, (k + 1) % n_folds) for pid in fold_pair_ids[j]
        ]
        folds.append({"train": train_ids, "valid": valid_ids, "test": test_ids})
    return folds


def assert_pair_integrity(split: dict[str, list[str]]) -> None:
    train, valid, test = set(split["train"]), set(split["valid"]), set(split["test"])
    assert not (train & valid), "pair_id leaked between train/valid"
    assert not (train & test), "pair_id leaked between train/test"
    assert not (valid & test), "pair_id leaked between valid/test"


def write_split(name: str, split: dict[str, list[str]], out_dir: Path = SPLITS_DIR) -> None:
    d = out_dir / name
    d.mkdir(parents=True, exist_ok=True)
    for part, ids in split.items():
        (d / f"{part}.json").write_text(json.dumps(sorted(ids)))