#!/usr/bin/env python3
"""
DPO Dataset — manifest-based preference pairs for HazardHead alignment.

Each sample is a (chosen, rejected) window pair where:
  chosen   = window where issuing an alert is CORRECT
             (ego_pos, TTA ∈ [1.5, 5.0]s  →  "timely_alert")
  rejected = window where issuing an alert is WRONG
             (too_early, too_late, safe_neg, non_ego)

The dataset returns raw PIL frames; the DPO trainer handles VLM tokenisation.
"""

from __future__ import annotations

import json
import logging
from pathlib import Path
from typing import Any, Dict, List, Optional

import torch
from PIL import Image
from torch.utils.data import Dataset

logger = logging.getLogger(__name__)

MAX_FRAMES = 8


# ─────────────────────────────────────────────────────────────────────────────
# Frame loader (mirrors SFT dataset)
# ─────────────────────────────────────────────────────────────────────────────

def _load_frame(src_dir: Path, frame_idx: int) -> Optional[Image.Image]:
    for fmt in ["{:03d}", "{:04d}", "{:05d}", "{:06d}", "{}"]:
        for ext in [".jpg", ".jpeg", ".png"]:
            p = src_dir / (fmt.format(frame_idx) + ext)
            if p.exists():
                try:
                    return Image.open(p).convert("RGB")
                except Exception:
                    pass
    return None


def _load_frames(source_dir: str, frame_indices: List[int]) -> List[Image.Image]:
    src = Path(source_dir)
    imgs = []
    for idx in frame_indices[:MAX_FRAMES]:
        img = _load_frame(src, idx)
        if img is not None:
            imgs.append(img)
    if not imgs:
        imgs = [Image.new("RGB", (384, 384), (64, 64, 64))]
    return imgs


# ─────────────────────────────────────────────────────────────────────────────
# DPODataset
# ─────────────────────────────────────────────────────────────────────────────

class DPODataset(Dataset):
    """
    Loads preference pairs from DPO pair manifests.

    Args
    ----
    manifests : list of paths to JSON pair manifests (as generated by make_dpo_pairs.py)
    split     : "train" or "val"
    debug     : if True, limit to debug_samples pairs
    debug_samples : number of pairs to use in debug mode
    """

    def __init__(
        self,
        manifests:     List[Path],
        split:         str = "train",
        debug:         bool = False,
        debug_samples: int = 64,
    ):
        self.split = split
        self.pairs: List[dict] = []

        for m in manifests:
            m = Path(m)
            if not m.exists():
                logger.warning(f"DPO manifest not found: {m}")
                continue
            with open(m) as f:
                data = json.load(f)
            p = data.get("pairs", [])
            self.pairs.extend(p)
            logger.info(f"Loaded {len(p)} pairs from {m.name}")

        if debug:
            self.pairs = self.pairs[:debug_samples]

        logger.info(
            f"DPODataset [{split}]: {len(self.pairs)} pairs "
            f"({sum(1 for p in self.pairs if p['pair_type']=='timing')} timing, "
            f"{sum(1 for p in self.pairs if p['pair_type']=='category')} category)"
        )

    def __len__(self) -> int:
        return len(self.pairs)

    def __getitem__(self, idx: int) -> Dict[str, Any]:
        pair = self.pairs[idx]
        c    = pair["chosen"]
        r    = pair["rejected"]

        chosen_images   = _load_frames(c["source_dir"], c["frame_indices"])
        rejected_images = _load_frames(r["source_dir"], r["frame_indices"])

        return {
            "pair_id":          pair["pair_id"],
            "video_id":         pair["video_id"],
            "source":           pair["source"],
            "pair_type":        pair["pair_type"],
            # chosen
            "chosen_images":    chosen_images,
            "chosen_tta":       float(c["tta_true"]),
            "chosen_label":     c["label"],
            "chosen_metadata":  c.get("metadata", {}),
            # rejected
            "rejected_images":  rejected_images,
            "rejected_tta":     float(r["tta_true"]),
            "rejected_label":   r["label"],
            "rejected_metadata":r.get("metadata", {}),
        }


# ─────────────────────────────────────────────────────────────────────────────
# Collate
# ─────────────────────────────────────────────────────────────────────────────

def dpo_collate_fn(batch: List[Dict[str, Any]]) -> Dict[str, Any]:
    return {
        "pair_ids":           [b["pair_id"]          for b in batch],
        "video_ids":          [b["video_id"]          for b in batch],
        "sources":            [b["source"]            for b in batch],
        "pair_types":         [b["pair_type"]         for b in batch],
        # chosen
        "chosen_images":      [b["chosen_images"]     for b in batch],
        "chosen_ttas":        torch.tensor([b["chosen_tta"]    for b in batch], dtype=torch.float32),
        "chosen_labels":      [b["chosen_label"]      for b in batch],
        "chosen_metadata":    [b["chosen_metadata"]   for b in batch],
        # rejected
        "rejected_images":    [b["rejected_images"]   for b in batch],
        "rejected_ttas":      torch.tensor([b["rejected_tta"]  for b in batch], dtype=torch.float32),
        "rejected_labels":    [b["rejected_label"]    for b in batch],
        "rejected_metadata":  [b["rejected_metadata"] for b in batch],
    }