#!/usr/bin/env python3
"""
Generate deterministic train/val split manifests for SFT.

Rules
-----
- NEXAR train positive + negative → 85/15 hash split by video_id
- DADA positive → 85/15 hash split; exclude acc_frame >= num_frames
- DADA non-ego  → 85/15 hash split; accident_time kept for sampling density only
- DADA negative (3 videos) → all train
- FOLDER ASSIGNMENT is the source of truth (ignore stale `accident` boolean)
- All timestamps are 20 Hz frame indices

Outputs (in --out_dir)
----------------------
  nexar_train.json, nexar_val.json
  dada_pos_train.json, dada_pos_val.json
  dada_noneego_train.json, dada_noneego_val.json
  dada_neg_train.json
  nexar_test_public.json  (diagnostic only — NOT used for checkpoint selection)
"""

import argparse
import hashlib
import json
import logging
from datetime import date
from pathlib import Path
from typing import Any, Dict, List, Optional

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger(__name__)

FRAME_RATE_HZ = 20


# ── helpers ──────────────────────────────────────────────────────────────────

def _hash_split(video_id: str, val_pct: int = 15) -> str:
    """Deterministic split: val if MD5(video_id) % 100 < val_pct."""
    h = int(hashlib.md5(video_id.encode()).hexdigest(), 16)
    return "val" if (h % 100) < val_pct else "train"


def _load_ann(path: Path) -> Optional[Dict[str, Any]]:
    try:
        return json.loads(path.read_text(encoding="utf-8", errors="ignore").lstrip("\ufeff"))
    except Exception as e:
        logger.debug(f"Failed to load {path}: {e}")
        return None


def _safe_int(x: Any) -> Optional[int]:
    if x is None:
        return None
    try:
        return int(float(str(x).strip()))
    except Exception:
        return None


def _count_frames(vd: Path) -> int:
    return sum(1 for f in vd.iterdir() if f.suffix.lower() in {".jpg", ".jpeg", ".png"})


def _entry(
    video_id: str,
    source: str,
    category: str,       # "ego_positive" | "non_ego" | "safe_neg"
    source_dir: Path,
    num_frames: int,
    accident_frame: Optional[int],   # 20Hz frame index
    risky_frame: Optional[int],      # 20Hz frame index
    metadata: Dict[str, Any],
) -> Dict[str, Any]:
    return {
        "video_id": video_id,
        "source": source,
        "category": category,
        "source_dir": str(source_dir.resolve()),
        "num_frames": num_frames,
        "accident_frame": accident_frame,
        "risky_frame": risky_frame,
        "metadata": metadata,
    }


def _meta(ann: Dict[str, Any]) -> Dict[str, Any]:
    return {
        "accident_type": ann.get("accident_type", ""),
        "weather": ann.get("weather", ""),
        "road_type": ann.get("road_type", ""),
        "car_speed": ann.get("car_speed", ""),
        "time_of_day": ann.get("time_of_day", ""),
    }


# ── NEXAR ─────────────────────────────────────────────────────────────────────

def process_nexar_train(nexar_root: Path, val_pct: int) -> Dict[str, List]:
    splits: Dict[str, List] = {"train": [], "val": []}
    train_dir = nexar_root / "train"
    if not train_dir.exists():
        logger.warning(f"NEXAR train not found: {train_dir}")
        return splits

    for cat_folder, cat_label in [("positive", "ego_positive"), ("negative", "safe_neg")]:
        cat_dir = train_dir / cat_folder
        if not cat_dir.exists():
            continue
        ok = skip = 0
        for vd in sorted(cat_dir.iterdir()):
            if not vd.is_dir():
                continue
            ann = _load_ann(vd / "annotation.json")
            if ann is None:
                continue
            nf = _count_frames(vd)
            if nf == 0:
                continue

            video_id = f"nexar_{vd.name}"

            # NEXAR train clips: use accident_time directly (no _local suffix in train)
            if cat_label == "ego_positive":
                acc = _safe_int(ann.get("accident_time_local") or ann.get("accident_time"))
                rsk = _safe_int(ann.get("risky_time_local") or ann.get("risky_time"))
                if acc is None or acc >= nf:
                    skip += 1
                    continue
            else:
                acc = rsk = None

            e = _entry(video_id, "nexar", cat_label, vd, nf, acc, rsk, _meta(ann))
            splits[_hash_split(video_id, val_pct)].append(e)
            ok += 1
        logger.info(f"  NEXAR train/{cat_folder}: {ok} ok, {skip} skipped")

    return splits


def process_nexar_test_public(nexar_root: Path) -> List:
    """Diagnostic only — NOT used for checkpoint selection."""
    entries = []
    test_dir = nexar_root / "test-public"
    if not test_dir.exists():
        return entries

    for cat_folder, cat_label in [("positive", "ego_positive"), ("negative", "safe_neg")]:
        cat_dir = test_dir / cat_folder
        if not cat_dir.exists():
            continue
        for vd in sorted(cat_dir.iterdir()):
            if not vd.is_dir():
                continue
            ann = _load_ann(vd / "annotation.json")
            if ann is None:
                continue
            nf = _count_frames(vd)
            if nf == 0:
                continue

            video_id = f"nexar_{vd.name}"
            if cat_label == "ego_positive":
                acc = _safe_int(ann.get("accident_time_local") or ann.get("accident_time"))
                rsk = _safe_int(ann.get("risky_time_local") or ann.get("risky_time"))
                if acc is None or acc >= nf:
                    continue
            else:
                acc = rsk = None

            entries.append(_entry(video_id, "nexar", cat_label, vd, nf, acc, rsk, _meta(ann)))

    return entries


# ── DADA ──────────────────────────────────────────────────────────────────────

def process_dada_positive(dada_root: Path, val_pct: int) -> Dict[str, List]:
    splits: Dict[str, List] = {"train": [], "val": []}
    pos_dir = dada_root / "positive"
    if not pos_dir.exists():
        logger.warning(f"DADA positive not found: {pos_dir}")
        return splits

    ok = skip_nf = skip_acc = 0
    for vd in sorted(pos_dir.iterdir()):
        if not vd.is_dir():
            continue
        ann = _load_ann(vd / "annotation.json")
        if ann is None:
            continue
        nf = _count_frames(vd)
        if nf == 0:
            skip_nf += 1
            continue

        # FOLDER = source of truth; ignore stale accident boolean
        acc = _safe_int(ann.get("accident_time"))
        rsk = _safe_int(ann.get("risky_time"))

        if acc is None or acc >= nf:
            skip_acc += 1
            logger.debug(f"DADA pos skip {vd.name}: acc={acc}, nf={nf}")
            continue

        # risky_frame=0 is valid (risk from video start)
        if rsk is not None:
            rsk = max(0, rsk)

        video_id = f"dada_{vd.name}"
        e = _entry(video_id, "dada", "ego_positive", vd, nf, acc, rsk, _meta(ann))
        splits[_hash_split(video_id, val_pct)].append(e)
        ok += 1

    logger.info(f"  DADA positive: {ok} ok, {skip_acc} invalid acc_frame, {skip_nf} no frames")
    return splits


def process_dada_noneego(dada_root: Path, val_pct: int) -> Dict[str, List]:
    splits: Dict[str, List] = {"train": [], "val": []}
    ne_dir = dada_root / "non-ego"
    if not ne_dir.exists():
        logger.warning(f"DADA non-ego not found: {ne_dir}")
        return splits

    ok = 0
    for vd in sorted(ne_dir.iterdir()):
        if not vd.is_dir():
            continue
        ann = _load_ann(vd / "annotation.json")
        if ann is None:
            continue
        nf = _count_frames(vd)
        if nf == 0:
            continue

        # FOLDER = source of truth: non-ego
        # accident_time / risky_time kept ONLY for near-accident oversampling
        acc = _safe_int(ann.get("accident_time"))
        rsk = _safe_int(ann.get("risky_time"))

        # Clamp for safety (won't be used as training label)
        if acc is not None:
            acc = min(max(0, acc), nf - 1)
        if rsk is not None:
            rsk = min(max(0, rsk), nf - 1)

        video_id = f"dada_{vd.name}"
        e = _entry(video_id, "dada", "non_ego", vd, nf, acc, rsk, _meta(ann))
        splits[_hash_split(video_id, val_pct)].append(e)
        ok += 1

    logger.info(f"  DADA non-ego: {ok} total")
    return splits


def process_dada_negative(dada_root: Path) -> List:
    """All DADA negatives go to train (only 3 videos)."""
    entries = []
    neg_dir = dada_root / "negative"
    if not neg_dir.exists():
        return entries

    for vd in sorted(neg_dir.iterdir()):
        if not vd.is_dir():
            continue
        ann = _load_ann(vd / "annotation.json")
        if ann is None:
            continue
        nf = _count_frames(vd)
        if nf == 0:
            continue

        video_id = f"dada_{vd.name}"
        entries.append(_entry(video_id, "dada", "safe_neg", vd, nf, None, None, _meta(ann)))

    logger.info(f"  DADA negative (all train): {len(entries)}")
    return entries


# ── write ─────────────────────────────────────────────────────────────────────

def write_manifest(out_dir: Path, name: str, split: str, videos: List) -> Path:
    out_dir.mkdir(parents=True, exist_ok=True)
    manifest = {
        "name": name,
        "split": split,
        "generated_at": str(date.today()),
        "frame_rate_hz": FRAME_RATE_HZ,
        "num_videos": len(videos),
        "category_counts": {
            cat: sum(1 for v in videos if v["category"] == cat)
            for cat in ("ego_positive", "non_ego", "safe_neg")
        },
        "videos": videos,
    }
    path = out_dir / f"{name}.json"
    path.write_text(json.dumps(manifest, indent=2))
    logger.info(f"  → {path}  ({len(videos)} videos)")
    return path


# ── main ──────────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--nexar_root", default="PROJECT_ROOT/NEXAR_COLLISION/dataset")
    parser.add_argument("--dada_root",  default="PROJECT_ROOT/DADA-2000")
    parser.add_argument("--out_dir",    default="PROJECT_ROOT/data/sft_manifests")
    parser.add_argument("--val_pct",    type=int, default=15, help="Percent of videos in val (default 15)")
    args = parser.parse_args()

    nexar_root = Path(args.nexar_root)
    dada_root  = Path(args.dada_root)
    out_dir    = Path(args.out_dir)
    val_pct    = args.val_pct

    logger.info("=" * 60)
    logger.info("Generating SFT split manifests")
    logger.info(f"  NEXAR: {nexar_root}")
    logger.info(f"  DADA:  {dada_root}")
    logger.info(f"  Out:   {out_dir}")
    logger.info(f"  Val %: {val_pct}%")
    logger.info("=" * 60)

    # NEXAR
    logger.info("Processing NEXAR train...")
    nexar = process_nexar_train(nexar_root, val_pct)
    write_manifest(out_dir, "nexar_train", "train", nexar["train"])
    write_manifest(out_dir, "nexar_val",   "val",   nexar["val"])

    logger.info("Processing NEXAR test-public (diagnostic)...")
    nexar_test = process_nexar_test_public(nexar_root)
    write_manifest(out_dir, "nexar_test_public", "test_public", nexar_test)

    # DADA
    logger.info("Processing DADA positive...")
    dada_pos = process_dada_positive(dada_root, val_pct)
    write_manifest(out_dir, "dada_pos_train", "train", dada_pos["train"])
    write_manifest(out_dir, "dada_pos_val",   "val",   dada_pos["val"])

    logger.info("Processing DADA non-ego...")
    dada_ne = process_dada_noneego(dada_root, val_pct)
    write_manifest(out_dir, "dada_noneego_train", "train", dada_ne["train"])
    write_manifest(out_dir, "dada_noneego_val",   "val",   dada_ne["val"])

    logger.info("Processing DADA negative...")
    dada_neg = process_dada_negative(dada_root)
    write_manifest(out_dir, "dada_neg_train", "train", dada_neg)

    # Summary
    n_pos_tr  = (sum(1 for e in nexar["train"] if e["category"]=="ego_positive")
               + len(dada_pos["train"]))
    n_pos_val = (sum(1 for e in nexar["val"]   if e["category"]=="ego_positive")
               + len(dada_pos["val"]))
    n_ne_tr   = len(dada_ne["train"])
    n_ne_val  = len(dada_ne["val"])
    n_neg_tr  = (sum(1 for e in nexar["train"] if e["category"]=="safe_neg")
               + len(dada_neg))
    n_neg_val = sum(1 for e in nexar["val"] if e["category"]=="safe_neg")

    logger.info("")
    logger.info("=" * 60)
    logger.info("SUMMARY")
    logger.info(f"  TRAIN ego_positive : {n_pos_tr}")
    logger.info(f"  TRAIN non_ego      : {n_ne_tr}")
    logger.info(f"  TRAIN safe_neg     : {n_neg_tr}")
    logger.info(f"  ---")
    logger.info(f"  VAL   ego_positive : {n_pos_val}  ← checkpoint selection")
    logger.info(f"  VAL   non_ego      : {n_ne_val}   ← false-alert monitoring")
    logger.info(f"  VAL   safe_neg     : {n_neg_val}")
    logger.info(f"  TEST  (nexar only) : {len(nexar_test)}  (diagnostic, NOT for ckpt sel.)")
    logger.info("=" * 60)


if __name__ == "__main__":
    main()