#!/usr/bin/env python3
"""
PretrainDataset for both Stage-A (single frame) and Stage-B (8-frame sequence).
"""

import json
import random
from pathlib import Path
from typing import List, Dict, Any

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


class PretrainDataset(Dataset):
    """
    Loads a JSON file produced by prepare_stage_a.py or prepare_stage_b.py.

    Stage-A samples have `image_path` (str) → single-frame.
    Stage-B samples have `frame_paths` (list[str]) → multi-frame sequence.
    """

    def __init__(self, json_path: str, split: str = "train"):
        self.split = split
        data = json.loads(Path(json_path).read_text(encoding="utf-8"))
        self.samples = [s for s in data if s.get("split", split) == split]
        if not self.samples:
            # Accept any split (for files that don't tag split)
            self.samples = data

        if split == "train":
            random.shuffle(self.samples)

        # Detect stage from first sample
        self.is_stage_b = "frame_paths" in self.samples[0]

        print(f"PretrainDataset [{split}]: {len(self.samples)} samples  "
              f"({'Stage-B multi-frame' if self.is_stage_b else 'Stage-A single-frame'})")

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx: int) -> Dict[str, Any]:
        s = self.samples[idx]
        task   = s["task"]
        prompt = s["prompt"]
        label  = s["label"]

        if self.is_stage_b:
            # Multi-frame: load each frame as PIL
            frames = []
            for fp in s["frame_paths"]:
                try:
                    frames.append(Image.open(fp).convert("RGB"))
                except Exception:
                    pass
            if not frames:
                # Fallback: white image
                frames = [Image.new("RGB", (224, 224), color=(128, 128, 128))]
        else:
            # Single-frame
            try:
                img = Image.open(s["image_path"]).convert("RGB")
            except Exception:
                img = Image.new("RGB", (224, 224), color=(128, 128, 128))
            frames = [img]

        return {
            "frames": frames,   # List[PIL.Image]
            "prompt": prompt,
            "label":  label,
            "task":   task,
        }


def collate_fn(batch: List[Dict]) -> Dict:
    """
    Collate individual samples into a batch dict.
    Keeps frames as-is (list of lists of PIL images) for the processor.
    """
    return {
        "frames":  [item["frames"]  for item in batch],
        "prompts": [item["prompt"]  for item in batch],
        "labels":  [item["label"]   for item in batch],
        "tasks":   [item["task"]    for item in batch],
    }