"""Phase 1 §3.3 CLI: reads data/processed/corpus.parquet (+ the SecVulEval
GitHub-resolved commit dates, if available) and writes all three split
regimes under data/splits/.
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path

import pandas as pd

from src.data.build_splits import (
    assert_pair_integrity,
    build_chronological_split,
    build_random_split,
    build_repo_disjoint_folds,
    _pair_table,
    write_split,
)
from src.utils.seeding import set_seed

SECVULEVAL_DATES_CACHE = Path("data/processed/secvuleval_commit_dates.jsonl")


def _merge_secvuleval_dates(df: pd.DataFrame, cache_path: Path = SECVULEVAL_DATES_CACHE) -> pd.DataFrame:
    """Fills in `commit_date` for secvuleval rows from the GitHub-resolved
    cache (src/data/fetch_commit_dates.py), where PrimeVul's Big-Vul join
    didn't already provide one. Does not overwrite an existing date."""
    if not cache_path.exists():
        print(f"no commit-date cache at {cache_path}; secvuleval rows stay undated")
        return df
    resolved: dict[str, str] = {}
    with open(cache_path) as f:
        for line in f:
            rec = json.loads(line)
            if rec.get("date"):
                resolved[rec["commit_id"]] = rec["date"]
    print(f"loaded {len(resolved)} resolved commit dates from {cache_path}")

    df = df.copy()
    needs_date = df["commit_date"].isna()
    df.loc[needs_date, "commit_date"] = df.loc[needs_date, "commit_id"].map(resolved)
    return df


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--corpus", default="data/processed/corpus.parquet")
    parser.add_argument("--chrono-train-before", default="2022-01-01")
    parser.add_argument("--chrono-valid-before", default="2023-01-01")
    parser.add_argument("--n-folds", type=int, default=5)
    parser.add_argument("--results-dir", default="results/phase1")
    args = parser.parse_args()
    set_seed(args.seed)

    results_dir = Path(args.results_dir)
    results_dir.mkdir(parents=True, exist_ok=True)

    df = pd.read_parquet(args.corpus)
    df = _merge_secvuleval_dates(df)
    pairs = _pair_table(df)
    print(f"{len(pairs)} pairs total; date coverage now {pairs['commit_date'].notna().mean():.1%}")

    random_split = build_random_split(pairs, seed=args.seed)
    assert_pair_integrity(random_split)
    write_split("random", random_split)
    print(f"random split: {[(k, len(v)) for k, v in random_split.items()]}")

    chrono_split, chrono_stats = build_chronological_split(
        pairs, args.chrono_train_before, args.chrono_valid_before
    )
    assert_pair_integrity(chrono_split)
    write_split("chronological", chrono_split)
    print(f"chronological split: {chrono_stats}")

    folds = build_repo_disjoint_folds(pairs, n_folds=args.n_folds, seed=args.seed)
    for i, fold in enumerate(folds):
        assert_pair_integrity(fold)
        write_split(f"repo_disjoint_fold{i}", fold)
    print(f"repo-disjoint: {args.n_folds} folds written")

    summary = {
        "n_pairs_total": len(pairs),
        "date_coverage_after_secvuleval_merge": float(pairs["commit_date"].notna().mean()),
        "random_split_sizes": {k: len(v) for k, v in random_split.items()},
        "chronological_split_stats": chrono_stats,
        "repo_disjoint_fold_sizes": [
            {k: len(v) for k, v in fold.items()} for fold in folds
        ],
    }
    (results_dir / "split_summary.json").write_text(json.dumps(summary, indent=2))
    print(json.dumps(summary, indent=2))


if __name__ == "__main__":
    main()