# Adopted from https://github.com/guandeh17/Self-Forcing
# SPDX-License-Identifier: Apache-2.0
from torch.utils.data import Dataset
import numpy as np
import torch
import random
import json
from pathlib import Path
from PIL import Image
import os
import subprocess
import time
import warnings
import torchvision.transforms as transforms
import torchvision.transforms.functional as F

try:
    import decord
except ModuleNotFoundError:
    decord = None

DEFAULT_SCENE_CUT_PREFIX = "The scene transitions. "


class TextDataset(Dataset):
    def __init__(self, prompt_path, extended_prompt_path=None):
        with open(prompt_path, encoding="utf-8") as f:
            self.prompt_list = [line.rstrip() for line in f]

        if extended_prompt_path is not None:
            with open(extended_prompt_path, encoding="utf-8") as f:
                self.extended_prompt_list = [line.rstrip() for line in f]
            assert len(self.extended_prompt_list) == len(self.prompt_list)
        else:
            self.extended_prompt_list = None

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

    def __getitem__(self, idx):
        batch = {
            "prompts": self.prompt_list[idx],
            "idx": idx,
        }
        if self.extended_prompt_list is not None:
            batch["extended_prompts"] = self.extended_prompt_list[idx]
        return batch


class MultiTextDataset(Dataset):
    """Dataset for multi-segment prompts stored in a JSONL file.

    Each line is a JSON object, e.g.
        {"prompts": ["a cat", "a dog", "a bird"]}

    Args
    ----
    prompt_path : str
        Path to the JSONL file
    field       : str
        Name of the list-of-strings field, default "prompts"
    cache_dir   : str | None
        ``cache_dir`` passed to HF Datasets (optional)
    """

    def __init__(self, prompt_path: str, field: str = "prompts", cache_dir: str | None = None):
        try:
            import datasets
        except ModuleNotFoundError as exc:
            raise ModuleNotFoundError(
                "The 'datasets' package is required for MultiTextDataset. "
                "Use MultiTextConcatDataset for plain txt/json-caption directories "
                "or install datasets."
            ) from exc
        self.ds = datasets.load_dataset(
            "json",
            data_files=prompt_path,
            split="train",
            cache_dir=cache_dir,
            streaming=False, 
        )

        assert len(self.ds) > 0, "JSONL is empty"
        assert field in self.ds.column_names, f"Missing field '{field}'"

        seg_len = len(self.ds[0][field])
        for i, ex in enumerate(self.ds):
            val = ex[field]
            assert isinstance(val, list), f"Line {i} field '{field}' is not a list"
            assert len(val) == seg_len,  f"Line {i} list length mismatch"

        self.field = field

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

    def __getitem__(self, idx: int):
        return {
            "idx": idx,
            "prompts_list": self.ds[idx][self.field],  # List[str]
        }


class MultiTextConcatDataset(Dataset):
    """Text-only dataset for multi-shot training and inference.

    Supports two input modes:

    **txt file** — each line is one caption. Each sample uses the caption at
    index ``idx``, repeated ``num_blocks`` times (single-shot, no scene cut
    prefix).

    **directory** — reads ``caption/<subfolder>/*.json`` files (no video dir
    needed). Shot durations are resolved with a three-level fallback:

    1. ``shot_durations.txt`` in the caption subfolder (per-sample override)
    2. ``chunks_per_shot`` from config (global fixed repeat)
    3. Even distribution across all available captions

    Scene cut prefix is prepended at shot boundaries (first block of each
    shot except shot 0). Output is always exactly ``num_blocks`` prompts:
    truncated if too many, padded with the last caption if too few.
    """

    def __init__(
        self,
        data_path: str,
        num_blocks: int,
        chunks_per_shot: int = 0,
        scene_cut_prefix: str = DEFAULT_SCENE_CUT_PREFIX,
        caption_field: str = "caption",
        deterministic: bool = False,
    ):
        self.num_blocks = num_blocks
        self.chunks_per_shot = chunks_per_shot
        self.scene_cut_prefix = scene_cut_prefix
        self.caption_field = caption_field
        self.deterministic = deterministic

        path = Path(data_path)
        if data_path.endswith(".txt") or path.is_file():
            self._mode = "txt"
            with open(data_path, encoding="utf-8") as f:
                self._prompts = [line.rstrip() for line in f if line.strip()]
            assert len(self._prompts) > 0, f"No prompts found in {data_path}"
        else:
            self._mode = "dir"
            self._caption_dir = path / "caption" if (path / "caption").is_dir() else path
            self._folders = sorted([d for d in self._caption_dir.iterdir() if d.is_dir()])
            assert len(self._folders) > 0, (
                f"No caption subfolders found in {self._caption_dir}"
            )

    def __len__(self):
        if self._mode == "txt":
            return len(self._prompts)
        return len(self._folders)

    def __getitem__(self, idx):
        if self._mode == "txt":
            return self._get_txt_item(idx)
        return self._get_dir_item(idx)

    # ------------------------------------------------------------------
    # txt mode
    # ------------------------------------------------------------------

    def _get_txt_item(self, idx):
        caption = self._prompts[idx % len(self._prompts)]
        return {
            "prompts": [caption] * self.num_blocks,
            "idx": idx,
        }

    # ------------------------------------------------------------------
    # directory mode
    # ------------------------------------------------------------------

    def _get_dir_item(self, idx):
        folder = self._folders[idx % len(self._folders)]
        raw_captions = self._load_captions_from_folder(folder)
        if not raw_captions:
            raw_captions = [""]

        shot_durations = self._resolve_shot_durations(folder, len(raw_captions))
        prompts = self._apply_shot_durations(raw_captions, shot_durations)

        # Ensure exactly num_blocks prompts
        if len(prompts) > self.num_blocks:
            prompts = prompts[: self.num_blocks]
        elif len(prompts) < self.num_blocks:
            last = prompts[-1] if prompts else ""
            prompts.extend([last] * (self.num_blocks - len(prompts)))

        return {
            "prompts": prompts,
            "idx": idx,
        }

    def _load_captions_from_folder(self, folder: Path):
        json_files = sorted(
            [f for f in folder.glob("*.json") if f.name != "global.json"],
            key=lambda p: (p.stem.isdigit(), int(p.stem) if p.stem.isdigit() else 0, p.stem),
        )
        captions = []
        for jf in json_files:
            try:
                with open(jf, "r", encoding="utf-8") as f:
                    data = json.load(f)
                    captions.append(data.get(self.caption_field, ""))
            except Exception:
                captions.append("")
        return captions

    # ------------------------------------------------------------------
    # shot duration helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _load_shot_durations(folder: Path):
        txt_path = folder / "shot_durations.txt"
        if not txt_path.exists():
            return None
        try:
            with open(txt_path, "r") as f:
                content = f.read().strip()
            parts = content.replace(",", " ").split()
            durations = [int(x) for x in parts if x.strip()]
            return durations if durations else None
        except Exception:
            return None

    def _resolve_shot_durations(self, folder: Path, num_captions: int):
        durations = self._load_shot_durations(folder)
        if durations is not None:
            return durations[:num_captions]
        if self.chunks_per_shot > 0:
            return [self.chunks_per_shot] * num_captions
        return self._even_durations(num_captions)

    def _even_durations(self, num_shots: int):
        total = self.num_blocks
        base, extra = divmod(total, num_shots)
        return [base + (1 if i < extra else 0) for i in range(num_shots)]

    def _apply_shot_durations(self, raw_captions, shot_durations):
        target = self.num_blocks
        clamped: list[int] = []
        remaining = target
        for d in shot_durations:
            if remaining <= 0:
                break
            take = min(d, remaining)
            clamped.append(take)
            remaining -= take
        if remaining > 0 and clamped:
            clamped[-1] += remaining

        prompts: list[str] = []
        for shot_idx, (caption, duration) in enumerate(zip(raw_captions, clamped)):
            for block_in_shot in range(duration):
                if shot_idx > 0 and block_in_shot == 0 and self.scene_cut_prefix:
                    prompts.append(self.scene_cut_prefix + caption)
                else:
                    prompts.append(caption)
        return prompts


class MultiVideoConcatDataset(Dataset):
    """Dataset that concatenates multiple videos from a folder into a fixed-length video.
    
    Each item consists of multiple video segments concatenated together:
    - First segment: first_chunk_frames frames (chunk)
    - Subsequent segments: subsequent_chunk_frames frames each (chunk)
    - Total: total_frames frames (first_chunk_frames + subsequent_chunk_frames*num_subsequent_segments = total_frames)
    
    Videos are sampled preserving original duration, and if a video doesn't have enough frames,
    it moves to the next video. If a video has enough frames, it can be sampled repeatedly.
    """
    def __init__(
        self,
        data_dir,
        video_size,
        total_frames,
        target_fps=16,
        video_extensions=('.mp4', '.avi', '.mov', '.mkv', '.webm'),
        caption_field='caption',
        filter_invalid_folders=False,
        deterministic: bool = False,
        num_frame_per_block=8,
        temporal_compression_ratio=4,
        allow_padding: bool = False,
        min_latent_frames: int = 0,
        single_video_only: bool = False,
        independent_first_frame: bool = False,
        return_image: bool = False,
        max_chunks_per_shot: int = 0,
        scene_cut_prefix: str = DEFAULT_SCENE_CUT_PREFIX,
        sample_warning_seconds: float = 60.0,
        sample_warning_interval_seconds: float = 60.0,
    ):
        self.root_dir = Path(data_dir)
        self.data_dir = self.root_dir / "video"
        self.caption_dir = self.root_dir / "caption"
        self.video_size = video_size
        self.total_frames = total_frames

        total_latent_frames = 1 + (total_frames - 1) // temporal_compression_ratio
        separate_first_latent = (
            independent_first_frame
            and total_latent_frames % num_frame_per_block != 0
        )
        if separate_first_latent:
            assert (total_latent_frames - 1) % num_frame_per_block == 0, (
                f"total latent frames ({total_latent_frames}) must be divisible by "
                f"num_frame_per_block ({num_frame_per_block}) or equal to "
                f"1 + N * num_frame_per_block when independent_first_frame=True"
            )
        first_chunk_latent_frames = (
            num_frame_per_block + 1 if separate_first_latent else num_frame_per_block
        )
        first_chunk_frames = 1 + (first_chunk_latent_frames - 1) * temporal_compression_ratio
        subsequent_chunk_frames = num_frame_per_block * temporal_compression_ratio

        self.first_chunk_frames = first_chunk_frames
        self.subsequent_chunk_frames = subsequent_chunk_frames
        self.target_fps = target_fps
        self.caption_field = caption_field
        self.video_extensions = video_extensions
        self.filter_invalid_folders = filter_invalid_folders
        self.deterministic = deterministic
        self.allow_padding = allow_padding
        self.num_frame_per_block = num_frame_per_block
        self.independent_first_frame = independent_first_frame
        self.first_chunk_latent_frames = first_chunk_latent_frames
        self.return_image = return_image
        if min_latent_frames > 0:
            assert min_latent_frames % num_frame_per_block == 0, (
                f"min_latent_frames ({min_latent_frames}) must be a multiple of "
                f"num_frame_per_block ({num_frame_per_block})"
            )
        self.min_latent_frames = min_latent_frames
        self.single_video_only = single_video_only
        self.max_chunks_per_shot = max_chunks_per_shot
        self.scene_cut_prefix = scene_cut_prefix
        self.sample_warning_seconds = float(sample_warning_seconds or 0.0)
        self.sample_warning_interval_seconds = float(sample_warning_interval_seconds or 0.0)
        
        remaining_frames = total_frames - first_chunk_frames
        self.num_subsequent_segments = remaining_frames // subsequent_chunk_frames
        self.total_segments = 1 + self.num_subsequent_segments
        
        assert total_frames == first_chunk_frames + self.num_subsequent_segments * subsequent_chunk_frames, \
            f"Total frames ({total_frames}) must equal first_chunk_frames ({first_chunk_frames}) + " \
            f"num_subsequent_segments ({self.num_subsequent_segments}) * subsequent_chunk_frames ({subsequent_chunk_frames})"
        
        if not self.data_dir.exists():
            raise ValueError(f"Video directory not found: {self.data_dir}")
        if not self.caption_dir.exists():
            raise ValueError(f"Caption directory not found: {self.caption_dir}")
        
        self.folders = [d for d in self.data_dir.iterdir() if d.is_dir()]
        if len(self.folders) == 0:
            raise ValueError(f"No subdirectories found in {self.data_dir}")

        # Optionally pre-filter folders with insufficient frames
        # Note: This can be slow for large datasets due to IO operations
        if self.filter_invalid_folders:
            print(f"[MultiVideoConcatDataset] Pre-filtering {len(self.folders)} folders for sufficient frames...")
            valid_folders = []
            skipped_folders = []
            for folder in self.folders:
                if self._check_folder_has_enough_frames(folder):
                    valid_folders.append(folder)
                else:
                    skipped_folders.append(folder.name)
            
            if len(skipped_folders) > 0:
                print(f"[MultiVideoConcatDataset] Skipped {len(skipped_folders)} folders due to insufficient frames: {skipped_folders}")
            
            self.folders = valid_folders
            
            if len(self.folders) == 0:
                raise ValueError(f"No folders with sufficient frames found in {self.data_dir}")
        
        # Setup transforms
        self.resize_transform = transforms.Resize(
            self.video_size, 
            antialias=True
        )
        self.normalize = transforms.Normalize(
            mean=[0.5, 0.5, 0.5],
            std=[0.5, 0.5, 0.5]
        )
        
        # Lazy caches: avoid repeated decord.VideoReader / filesystem IO
        self._video_info_cache = {}   # video_path -> (total_frames, fps)
        self._folder_files_cache = {} # folder_path -> list of video paths
        
        if decord is not None:
            decord.bridge.set_bridge('torch')
    
    def _get_caption_folder(self, folder_name):
        """Return the caption directory path for a given folder (sample)."""
        return self.caption_dir / folder_name

    def _load_caption(self, video_path, folder_name):
        """Load caption for a video file."""
        video_stem = video_path.stem
        caption_folder = self._get_caption_folder(folder_name)
        caption_path = caption_folder / f"{video_stem}.json"
        
        if caption_path.exists():
            try:
                with open(caption_path, 'r', encoding='utf-8') as f:
                    caption_data = json.load(f)
                    return caption_data.get(self.caption_field, "")
            except Exception:
                return ""
        return ""

    def _get_video_files_in_folder(self, folder_path):
        """Get sorted video files in a folder, keeping only those with a per-video caption (cached)."""
        key = str(folder_path)
        if key in self._folder_files_cache:
            return self._folder_files_cache[key]

        video_files = []
        for ext in self.video_extensions:
            video_files.extend(list(folder_path.glob(f'*{ext}')))
            video_files.extend(list(folder_path.glob(f'*{ext.upper()}')))
        
        def get_numeric_key(path):
            try:
                return int(path.stem)
            except ValueError:
                return float('inf')
        
        video_files.sort(key=get_numeric_key)

        folder_name = folder_path.name
        filtered_videos = []
        for video_path in video_files:
            caption = self._load_caption(video_path, folder_name)
            if caption is not None and caption != "":
                filtered_videos.append(video_path)

        self._folder_files_cache[key] = filtered_videos
        return filtered_videos
    
    def _check_folder_has_enough_frames(self, folder_path):
        """Check if a folder has enough total frames across all videos to complete all segments.
        
        This is a lenient check: we verify that the total available frames across all videos
        is sufficient for the required segments, assuming ideal sampling.
        """
        video_files = self._get_video_files_in_folder(folder_path)
        if len(video_files) == 0:
            return False
        
        # Calculate total available frames (in target_fps timebase)
        total_available_frames = 0
        for video_path in video_files:
            try:
                total_frames, original_fps = self._get_video_info(video_path)
                # Convert to target_fps timebase
                available_in_target_fps = total_frames * self.target_fps / original_fps
                total_available_frames += available_in_target_fps
            except Exception:
                # If we can't read a video, be conservative and skip this folder
                return False
        
        # Check if we have enough frames for all segments
        required_frames = self.total_frames
        return total_available_frames >= required_frames
    
    def _get_video_info(self, video_path):
        """Get video information without loading frames (cached)."""
        key = str(video_path)
        if key in self._video_info_cache:
            return self._video_info_cache[key]
        if decord is None:
            info = self._get_video_info_ffprobe(video_path)
            self._video_info_cache[key] = info
            return info
        try:
            vr = decord.VideoReader(key, width=self.video_size[1], height=self.video_size[0])
        except:
            vr = decord.VideoReader(key)
        info = (len(vr), vr.get_avg_fps())
        self._video_info_cache[key] = info
        return info

    @staticmethod
    def _parse_fps(value):
        if not value or value == "0/0":
            return 0.0
        if "/" in value:
            num, den = value.split("/", 1)
            den_f = float(den)
            return float(num) / den_f if den_f != 0 else 0.0
        return float(value)

    def _get_video_info_ffprobe(self, video_path):
        cmd = [
            "ffprobe",
            "-v",
            "error",
            "-select_streams",
            "v:0",
            "-count_frames",
            "-show_entries",
            "stream=nb_read_frames,nb_frames,avg_frame_rate,r_frame_rate,duration",
            "-of",
            "json",
            str(video_path),
        ]
        try:
            proc = subprocess.run(cmd, check=True, capture_output=True, text=True)
            stream = json.loads(proc.stdout)["streams"][0]
            fps = self._parse_fps(stream.get("avg_frame_rate")) or self._parse_fps(stream.get("r_frame_rate"))
            frames_raw = stream.get("nb_read_frames") or stream.get("nb_frames")
            if frames_raw and str(frames_raw).isdigit():
                total_frames = int(frames_raw)
            else:
                total_frames = int(round(float(stream.get("duration", 0.0)) * fps))
            if total_frames <= 0 or fps <= 0:
                raise ValueError(f"Could not infer frames/fps from ffprobe output for {video_path}")
            return total_frames, fps
        except (subprocess.CalledProcessError, KeyError, IndexError, ValueError, json.JSONDecodeError) as exc:
            raise RuntimeError(f"ffprobe failed to read video metadata for {video_path}: {exc}") from exc
    
    def _can_sample_from_position(self, total_frames, original_fps, num_frames, start_frame):
        """Check if we can sample num_frames starting from start_frame.
        
        Returns True if we can sample num_frames without exceeding video bounds.
        """
        if start_frame >= total_frames:
            return False
        
        # Calculate sampling interval: original_fps / target_fps
        sampling_interval = original_fps / self.target_fps
        
        # Calculate the last frame index we need (with rounding)
        # We need to check if we can get num_frames frames
        last_frame_needed = start_frame + (num_frames - 1) * sampling_interval
        
        # Account for rounding: the actual last frame index will be rounded
        # So we need some margin to ensure we don't exceed bounds
        return int(np.round(last_frame_needed)) < total_frames

    def _can_complete_all_segments_without_wrap(self, video_files, start_video_idx, start_frame):
        """Check if from (start_video_idx, start_frame) we can sample all segments
        without ever wrapping to the beginning (i.e. only use this video and later ones).
        When single_video_only is True, all segments must come from the same video.
        """
        vidx = start_video_idx
        start = start_frame

        # First segment
        total_frames, original_fps = self._get_video_info(video_files[vidx])
        if not self._can_sample_from_position(
            total_frames, original_fps, self.first_chunk_frames, start
        ):
            return False
        sampling_interval = original_fps / self.target_fps
        next_start = start + self.first_chunk_frames * sampling_interval
        if int(np.round(next_start)) >= total_frames:
            if self.single_video_only:
                return self.num_subsequent_segments == 0
            vidx += 1
            start = 0
        else:
            start = int(np.round(next_start))
        if vidx >= len(video_files):
            return False

        # Subsequent segments
        for seg_i in range(self.num_subsequent_segments):
            while vidx < len(video_files):
                total_frames, original_fps = self._get_video_info(video_files[vidx])
                if self._can_sample_from_position(
                    total_frames, original_fps, self.subsequent_chunk_frames, start
                ):
                    break
                if self.single_video_only:
                    return False
                vidx += 1
                start = 0
            if vidx >= len(video_files):
                return False
            sampling_interval = original_fps / self.target_fps
            next_start = start + self.subsequent_chunk_frames * sampling_interval
            if int(np.round(next_start)) >= total_frames:
                if self.single_video_only and seg_i < self.num_subsequent_segments - 1:
                    return False
                vidx += 1
                start = 0
            else:
                start = int(np.round(next_start))
        return True

    def _sample_random_start(self, video_files):
        """Sample a random (video_idx, start_frame) valid for the first segment,
        and from which we can complete ALL segments without wrapping to the start.
        Returns (video_idx, start_frame); falls back to (0, 0) if no valid start found.
        """
        candidates = []
        for video_idx, video_path in enumerate(video_files):
            total_frames, original_fps = self._get_video_info(video_path)
            sampling_interval = original_fps / self.target_fps
            last_needed = (self.first_chunk_frames - 1) * sampling_interval
            max_start = int(np.floor(total_frames - 1 - last_needed))
            if max_start < 0:
                continue
            step = max(1, max_start // 50)
            for start in range(0, max_start + 1, step):
                if not self._can_sample_from_position(
                    total_frames, original_fps, self.first_chunk_frames, start
                ):
                    continue
                if self._can_complete_all_segments_without_wrap(video_files, video_idx, start):
                    candidates.append((video_idx, start))
        if candidates:
            chosen = random.choice(candidates)
            return chosen
        return (0, 0)
    
    def _sample_frames_from_video(self, video_path, num_frames, start_frame=0):
        """Sample frames from a video preserving original duration.
        
        Args:
            video_path: Path to video file
            num_frames: Number of frames to sample
            start_frame: Starting frame index (for repeated sampling)
        
        Returns:
            tuple: (frames_tensor, total_frames_in_video, original_fps)
        """
        if decord is None:
            return self._sample_frames_from_video_ffmpeg(video_path, num_frames, start_frame)

        try:
            vr = decord.VideoReader(str(video_path), width=self.video_size[1], height=self.video_size[0])
        except:
            vr = decord.VideoReader(str(video_path))
        
        total_frames = len(vr)
        original_fps = vr.get_avg_fps()
        
        if total_frames == 0:
            raise ValueError(f"Video {video_path} has no frames")
        
        # Calculate frame sampling based on fps to preserve duration
        # Calculate sampling interval: original_fps / target_fps
        sampling_interval = original_fps / self.target_fps
        
        # Generate frame indices starting from start_frame
        indices = []
        current_frame = float(start_frame)
        for _ in range(num_frames):
            frame_idx = int(np.round(current_frame))
            frame_idx = min(frame_idx, total_frames - 1)
            indices.append(frame_idx)
            current_frame += sampling_interval
            if current_frame >= total_frames:
                # If we run out of frames, pad with the last frame
                remaining = num_frames - len(indices)
                indices.extend([total_frames - 1] * remaining)
                break
        
        indices = np.array(indices[:num_frames], dtype=np.int32)
        
        # Get video frames: shape (num_frames, height, width, 3)
        video_frames = vr.get_batch(indices).numpy()
        
        # Convert to tensor and permute to (num_frames, 3, height, width)
        video_tensor = torch.from_numpy(video_frames).permute(0, 3, 1, 2).contiguous()
        
        # Convert to float and normalize pixel values from [0, 255] to [0, 1]
        video_tensor = video_tensor.float() / 255.0
        
        # Resize if needed
        if video_tensor.shape[2] != self.video_size[0] or video_tensor.shape[3] != self.video_size[1]:
            resized_frames = []
            for i in range(video_tensor.shape[0]):
                resized_frames.append(self.resize_transform(video_tensor[i]))
            video_tensor = torch.stack(resized_frames, dim=0)
        
        # Apply normalization: (x - 0.5) / 0.5 -> range [-1, 1]
        video_tensor = self.normalize(video_tensor)
        video_tensor = video_tensor.to(torch.float16)
        
        return video_tensor, total_frames, original_fps

    def _sample_frames_from_video_ffmpeg(self, video_path, num_frames, start_frame=0):
        total_frames, original_fps = self._get_video_info(video_path)
        start_time = max(0.0, float(start_frame) / max(original_fps, 1e-6))
        height, width = self.video_size
        cmd = [
            "ffmpeg",
            "-v",
            "error",
            "-ss",
            f"{start_time:.6f}",
            "-i",
            str(video_path),
            "-vf",
            f"fps={self.target_fps},scale={width}:{height}",
            "-frames:v",
            str(num_frames),
            "-f",
            "rawvideo",
            "-pix_fmt",
            "rgb24",
            "pipe:1",
        ]
        try:
            proc = subprocess.run(cmd, check=True, capture_output=True)
        except subprocess.CalledProcessError as exc:
            stderr = exc.stderr.decode("utf-8", errors="replace") if exc.stderr else ""
            raise RuntimeError(f"ffmpeg failed to sample {video_path}: {stderr}") from exc

        frame_bytes = height * width * 3
        decoded_frames = len(proc.stdout) // frame_bytes
        if decoded_frames <= 0:
            raise RuntimeError(f"ffmpeg produced no frames for {video_path}")

        video_frames = np.frombuffer(proc.stdout[: decoded_frames * frame_bytes], dtype=np.uint8)
        video_frames = video_frames.reshape(decoded_frames, height, width, 3)
        video_tensor = torch.from_numpy(video_frames.copy()).permute(0, 3, 1, 2).contiguous()
        video_tensor = video_tensor.float() / 255.0
        if decoded_frames < num_frames:
            pad = video_tensor[-1:].repeat(num_frames - decoded_frames, 1, 1, 1)
            video_tensor = torch.cat([video_tensor, pad], dim=0)
        elif decoded_frames > num_frames:
            video_tensor = video_tensor[:num_frames]

        video_tensor = self.normalize(video_tensor)
        video_tensor = video_tensor.to(torch.float16)
        return video_tensor, total_frames, original_fps
    
    def __len__(self):
        return len(self.folders)
    
    @staticmethod
    def _sample_failure(reason):
        return False, reason

    def _try_get_item_from_folder(self, folder_idx, deterministic: bool = False):
        """Try to get item from a specific folder.

        Returns (True, result) on success and (False, reason) on failure.
        The reason is used by __getitem__ warnings so a long scan does not
        look like a silent hang when most folders are too short.
        """
        folder_path = self.folders[folder_idx]
        folder_name = folder_path.name
        
        video_files = self._get_video_files_in_folder(folder_path)
        
        if len(video_files) == 0:
            return self._sample_failure(f"{folder_name}: no videos with captions")
        
        # Collect all segments
        all_segments = []
        prompts_list = []
        
        try:
            # Start position
            if deterministic:
                current_video_idx, current_start_frame = 0, 0
            else:
                # Random start: don't always start from the first video, so we get more diversity
                current_video_idx, current_start_frame = self._sample_random_start(video_files)
            
            # Sample first segment (9 frames)
            video_path = video_files[current_video_idx]
            prev_seg_video_idx = current_video_idx
            total_frames, original_fps = self._get_video_info(video_path)

            # Ensure the current video is long enough for the first segment
            while not self._can_sample_from_position(
                total_frames, original_fps, self.first_chunk_frames, current_start_frame
            ):
                if self.single_video_only:
                    return self._sample_failure(
                        f"{folder_name}: first video is too short for the first chunk"
                    )
                current_video_idx += 1
                current_start_frame = 0
                if current_video_idx >= len(video_files):
                    return self._sample_failure(
                        f"{folder_name}: no video can provide the first chunk"
                    )
                video_path = video_files[current_video_idx]
                prev_seg_video_idx = current_video_idx
                total_frames, original_fps = self._get_video_info(video_path)

            # Sample first segment
            segment_frames, total_frames, original_fps = self._sample_frames_from_video(
                video_path, self.first_chunk_frames, current_start_frame
            )
            all_segments.append(segment_frames)
            
            prompt = self._load_caption(video_path, folder_name)
            prompts_list.append(prompt)
            
            # Update position for next sampling
            sampling_interval = original_fps / self.target_fps
            next_start_frame = current_start_frame + self.first_chunk_frames * sampling_interval
            
            # If we've exhausted this video, move to next
            if int(np.round(next_start_frame)) >= total_frames:
                if self.single_video_only:
                    if self.num_subsequent_segments > 0:
                        return self._sample_failure(
                            f"{folder_name}: single-video sample ends after first chunk"
                        )
                else:
                    current_video_idx += 1
                    current_start_frame = 0
            else:
                current_start_frame = int(np.round(next_start_frame))

            chunks_from_current_video = 1 if current_video_idx == prev_seg_video_idx else 0

            # Sample subsequent segments (12 frames each). No wrap: we only use videos from start onward.
            for seg_idx in range(self.num_subsequent_segments):
                if current_video_idx >= len(video_files):
                    if self.allow_padding:
                        break
                    return self._sample_failure(
                        f"{folder_name}: ran out of videos before all chunks were filled"
                    )

                # Force virtual scene cut if max_shot_chunks reached:
                # skip 1 second of video and treat the remainder as a new shot.
                forced_scene_cut = False
                if (self.max_chunks_per_shot > 0
                        and chunks_from_current_video >= self.max_chunks_per_shot):
                    vp = video_files[current_video_idx]
                    _, ofps = self._get_video_info(vp)
                    current_start_frame += int(np.round(ofps))
                    chunks_from_current_video = 0
                    forced_scene_cut = True

                video_path = video_files[current_video_idx]
                total_frames, original_fps = self._get_video_info(video_path)

                can_sample = True
                while not self._can_sample_from_position(
                    total_frames, original_fps, self.subsequent_chunk_frames, current_start_frame
                ):
                    if self.single_video_only:
                        can_sample = False
                        break
                    current_video_idx += 1
                    current_start_frame = 0
                    chunks_from_current_video = 0
                    if current_video_idx >= len(video_files):
                        can_sample = False
                        break
                    video_path = video_files[current_video_idx]
                    total_frames, original_fps = self._get_video_info(video_path)

                if not can_sample:
                    if self.allow_padding:
                        break
                    return self._sample_failure(
                        f"{folder_name}: remaining videos are too short for the next chunk"
                    )

                is_scene_cut = (current_video_idx != prev_seg_video_idx) or forced_scene_cut

                # Sample segment
                segment_frames, total_frames, original_fps = self._sample_frames_from_video(
                    video_path, self.subsequent_chunk_frames, current_start_frame
                )
                all_segments.append(segment_frames)
                
                prompt = self._load_caption(video_path, folder_name)
                if is_scene_cut and self.scene_cut_prefix:
                    prompt = self.scene_cut_prefix + prompt
                prompts_list.append(prompt)

                prev_seg_video_idx = current_video_idx
                chunks_from_current_video += 1

                # Update position for next sampling
                sampling_interval = original_fps / self.target_fps
                next_start_frame = current_start_frame + self.subsequent_chunk_frames * sampling_interval
                
                # If we've exhausted this video, move to next
                if int(np.round(next_start_frame)) >= total_frames:
                    if self.single_video_only:
                        if seg_idx < self.num_subsequent_segments - 1:
                            return self._sample_failure(
                                f"{folder_name}: single-video sample is too short"
                            )
                    else:
                        current_video_idx += 1
                        current_start_frame = 0
                        chunks_from_current_video = 0
                else:
                    current_start_frame = int(np.round(next_start_frame))

            num_filled_segments = len(all_segments)
            if num_filled_segments == 0:
                num_valid_latent_frames = 0
            else:
                num_valid_latent_frames = (
                    self.first_chunk_latent_frames
                    + (num_filled_segments - 1) * self.num_frame_per_block
                )

            # Reject if below minimum latent frame threshold
            if self.allow_padding and self.min_latent_frames > 0:
                if num_valid_latent_frames < self.min_latent_frames:
                    return self._sample_failure(
                        f"{folder_name}: only {num_valid_latent_frames} valid latent frames, "
                        f"below min_latent_frames={self.min_latent_frames}"
                    )

            if num_filled_segments < self.total_segments:
                last_prompt = prompts_list[-1] if prompts_list else ""
                prompts_list.extend([last_prompt] * (self.total_segments - num_filled_segments))
            
            # Concatenate all segments: (total_frames, 3, height, width)
            concatenated_video = torch.cat(all_segments, dim=0)
            
            # Ensure we have exactly total_frames
            if concatenated_video.shape[0] != self.total_frames:
                # Pad or trim if necessary
                if concatenated_video.shape[0] < self.total_frames:
                    # Pad with last frame
                    last_frame = concatenated_video[-1:].repeat(self.total_frames - concatenated_video.shape[0], 1, 1, 1)
                    concatenated_video = torch.cat([concatenated_video, last_frame], dim=0)
                else:
                    # Trim
                    concatenated_video = concatenated_video[:self.total_frames]
            
            result = {
                'frames': concatenated_video.permute(1, 0, 2, 3),
                'prompts': prompts_list,
                'idx': folder_idx
            }
            if self.return_image:
                result['image'] = concatenated_video[0]
            if self.allow_padding:
                result['num_valid_latent_frames'] = num_valid_latent_frames
            return True, result
        except Exception as exc:
            return self._sample_failure(f"{folder_name}: {type(exc).__name__}: {exc}")

    def __getitem__(self, idx):
        start_time = time.monotonic()
        last_warning_time = start_time
        attempts = 0
        last_failure = None

        def maybe_warn(folder_idx, failure_reason):
            nonlocal last_warning_time
            if self.sample_warning_seconds <= 0:
                return
            elapsed = time.monotonic() - start_time
            if elapsed < self.sample_warning_seconds:
                return
            if (
                self.sample_warning_interval_seconds > 0
                and time.monotonic() - last_warning_time < self.sample_warning_interval_seconds
            ):
                return
            last_warning_time = time.monotonic()
            folder_name = self.folders[folder_idx % len(self.folders)].name
            warnings.warn(
                "[MultiVideoConcatDataset] Still searching for a valid sample "
                f"after {elapsed:.1f}s and {attempts} folder attempts "
                f"(requested_idx={idx}, current_folder={folder_name}, "
                f"last_failure={failure_reason}). This usually means the dataset "
                f"does not contain enough video duration for total_frames={self.total_frames} "
                f"at target_fps={self.target_fps}. Consider reducing the training "
                "window, enabling allow_padding, lowering min_latent_frames, or "
                "pre-filtering invalid folders.",
                RuntimeWarning,
                stacklevel=2,
            )

        # First try the requested folder
        attempts += 1
        success, result = self._try_get_item_from_folder(idx, deterministic=self.deterministic)
        if success:
            return result
        last_failure = result
        maybe_warn(idx, last_failure)
        
        # If the requested folder fails, try other folders.
        # If any valid folder exists in the dataset, try to return data from it:
        # scan every other folder starting from idx + 1 and return the first
        # successful sample.
        num_folders = len(self.folders)
        for i in range(1, num_folders):
            alt_idx = (idx + i) % num_folders
            attempts += 1
            success, result = self._try_get_item_from_folder(
                alt_idx,
                deterministic=self.deterministic,
            )
            if success:
                return result
            last_failure = result
            maybe_warn(alt_idx, last_failure)
        
        # If all attempts fail, raise an error
        elapsed = time.monotonic() - start_time
        if self.sample_warning_seconds > 0 and elapsed >= self.sample_warning_seconds:
            warnings.warn(
                "[MultiVideoConcatDataset] No valid sample was found after "
                f"{elapsed:.1f}s and {attempts} folder attempts "
                f"(requested_idx={idx}, last_failure={last_failure}).",
                RuntimeWarning,
                stacklevel=2,
            )
        raise ValueError(
            f"Failed to sample valid data from folder {idx} and nearby folders. "
            f"This may indicate insufficient valid videos in the dataset. "
            f"Tried {attempts} folders in {elapsed:.1f}s. Last failure: {last_failure}"
        )


def cycle(dl):
    while True:
        for data in dl:
            yield data

def multi_video_collate_fn(batch):
    # batch is a length-B list of dictionaries returned by __getitem__.
    frames = torch.stack([b["frames"] for b in batch], dim=0)  # (B, T, C, H, W)

    # Keep prompts as one list per sample:
    # [[p0_seg0, p0_seg1, ...], [p1_seg0, ...], ...].
    prompts_list = [b["prompts"] for b in batch]          # List[List[str]]

    idx = torch.tensor([b["idx"] for b in batch], dtype=torch.long)

    result = {
        "frames": frames,
        "prompts": prompts_list,
        "idx": idx,
    }

    if "image" in batch[0]:
        result["image"] = torch.stack([b["image"] for b in batch], dim=0)

    if "num_valid_latent_frames" in batch[0]:
        result["num_valid_latent_frames"] = torch.tensor(
            [b["num_valid_latent_frames"] for b in batch], dtype=torch.long
        )

    return result


def eval_collate_fn(batch):
    """Collate for text-only datasets (no frames)."""
    prompts_list = [b["prompts"] for b in batch]
    idx = torch.tensor([b["idx"] for b in batch], dtype=torch.long)
    result = {
        "prompts": prompts_list,
        "idx": idx,
    }
    if "shot_durations" in batch[0]:
        result["shot_durations"] = [b["shot_durations"] for b in batch]
    return result