"""
Frame sampling utilities for video processing.

This module provides video iteration and frame sampling utilities for inference.
The core frame sampling logic is in badas.core.preprocessing (shared with training).

This module adds:
- Video file iteration utilities (OpenCV and Decord backends)
- Sliding window iteration for streaming inference
- Timestamp calculation utilities
"""

import cv2
import numpy as np
from typing import List, Tuple, Generator, Optional, Dict, Any

# Import shared frame sampling from core (used by both training and inference)
from badas.core.preprocessing import calculate_frame_indices

# Try to import decord, fall back to OpenCV if not available
try:
    from decord import VideoReader, cpu, gpu
    DECORD_AVAILABLE = True
except ImportError:
    DECORD_AVAILABLE = False
    VideoReader = None


def get_frame_interval(source_fps: float, target_fps: float) -> float:
    """
    Calculate the frame interval for sampling.

    IMPORTANT: Returns a FLOAT - do not truncate to int!

    Args:
        source_fps: FPS of the source video (e.g., 30.0)
        target_fps: Desired FPS for the model (e.g., 8.0)

    Returns:
        Float interval between sampled frames (e.g., 3.75)

    Example:
        >>> get_frame_interval(30.0, 8.0)
        3.75
    """
    if source_fps <= 0 or target_fps <= 0:
        raise ValueError(f"FPS must be positive: source={source_fps}, target={target_fps}")

    return source_fps / target_fps


# Note: calculate_frame_indices is imported from badas.core.preprocessing
# This ensures training and inference use the same frame sampling logic.


def calculate_timestamp(frame_index: int, target_fps: float) -> float:
    """
    Calculate the timestamp in seconds for a given sampled frame index.

    Args:
        frame_index: The frame index in the resampled video (at target_fps)
        target_fps: The FPS of the resampled video

    Returns:
        Timestamp in seconds
    """
    return frame_index / target_fps


def expected_sampled_frames(
    total_frames: int,
    source_fps: float,
    target_fps: float
) -> int:
    """
    Calculate how many frames will be sampled from a video.

    Args:
        total_frames: Total frames in source video
        source_fps: Source video FPS
        target_fps: Target sampling FPS

    Returns:
        Number of frames that will be sampled
    """
    return len(calculate_frame_indices(total_frames, source_fps, target_fps))


def get_video_info(video_path: str) -> Dict[str, Any]:
    """
    Get video information.

    Args:
        video_path: Path to the video file

    Returns:
        Dictionary with video info:
        - fps: Frames per second
        - total_frames: Total number of frames
        - duration: Duration in seconds
        - width: Frame width
        - height: Frame height
    """
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        raise ValueError(f"Cannot open video: {video_path}")

    try:
        fps = cap.get(cv2.CAP_PROP_FPS)
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        duration = total_frames / fps if fps > 0 else 0

        return {
            'fps': fps,
            'total_frames': total_frames,
            'duration': duration,
            'width': width,
            'height': height
        }
    finally:
        cap.release()


def iter_sampled_frames(
    video_path: str,
    target_fps: float,
    resize: Optional[Tuple[int, int]] = None,
    convert_rgb: bool = True
) -> Generator[Tuple[int, int, np.ndarray], None, None]:
    """
    Generator that yields sampled frames from a video at the target FPS.

    Uses seek-based sampling (good for random access).

    Args:
        video_path: Path to the video file
        target_fps: Desired FPS for sampling
        resize: Optional (width, height) to resize frames
        convert_rgb: Whether to convert BGR to RGB (default True)

    Yields:
        Tuple of (original_frame_idx, sampled_frame_idx, frame)
        - original_frame_idx: Index in the original video
        - sampled_frame_idx: Index in the resampled sequence (0, 1, 2, ...)
        - frame: The frame as numpy array

    Example:
        for orig_idx, sampled_idx, frame in iter_sampled_frames("video.mp4", 8.0):
            print(f"Sampled frame {sampled_idx} from original frame {orig_idx}")
    """
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        raise ValueError(f"Cannot open video: {video_path}")

    try:
        source_fps = cap.get(cv2.CAP_PROP_FPS)
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

        frame_interval = get_frame_interval(source_fps, target_fps)

        current_pos = 0.0
        sampled_idx = 0

        while current_pos < total_frames:
            target_frame = int(current_pos)

            # Seek to the target frame
            cap.set(cv2.CAP_PROP_POS_FRAMES, target_frame)
            ret, frame = cap.read()

            if not ret:
                break

            # Process frame
            if resize:
                frame = cv2.resize(frame, resize)
            if convert_rgb:
                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

            yield (target_frame, sampled_idx, frame)

            current_pos += frame_interval
            sampled_idx += 1

    finally:
        cap.release()


def iter_sampled_frames_sequential(
    video_path: str,
    target_fps: float,
    resize: Optional[Tuple[int, int]] = None,
    convert_rgb: bool = True
) -> Generator[Tuple[int, int, np.ndarray], None, None]:
    """
    Generator that yields sampled frames by reading sequentially (faster for most videos).

    Instead of seeking to each frame, reads all frames and yields only the sampled ones.
    This is typically faster because sequential reads are more efficient than seeks.

    Args:
        video_path: Path to the video file
        target_fps: Desired FPS for sampling
        resize: Optional (width, height) to resize frames
        convert_rgb: Whether to convert BGR to RGB (default True)

    Yields:
        Tuple of (original_frame_idx, sampled_frame_idx, frame)
    """
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        raise ValueError(f"Cannot open video: {video_path}")

    try:
        source_fps = cap.get(cv2.CAP_PROP_FPS)
        frame_interval = get_frame_interval(source_fps, target_fps)

        next_sample_pos = 0.0
        frame_idx = 0
        sampled_idx = 0

        while True:
            ret, frame = cap.read()
            if not ret:
                break

            # Check if this frame should be sampled using floating-point comparison
            if frame_idx >= int(next_sample_pos):
                # Process frame
                if resize:
                    frame = cv2.resize(frame, resize)
                if convert_rgb:
                    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

                yield (frame_idx, sampled_idx, frame)

                next_sample_pos += frame_interval
                sampled_idx += 1

            frame_idx += 1

    finally:
        cap.release()


class FrameSampler:
    """
    Stateful frame sampler for streaming video processing.

    Maintains floating-point accumulator for accurate frame sampling
    when processing frames one at a time (e.g., from a live stream).

    Example:
        sampler = FrameSampler(source_fps=30.0, target_fps=8.0)

        for frame in video_stream:
            should_sample, sampled_idx = sampler.should_sample()
            if should_sample:
                process_frame(frame, sampled_idx)
            sampler.advance()
    """

    def __init__(self, source_fps: float, target_fps: float):
        """
        Initialize the frame sampler.

        Args:
            source_fps: FPS of the source video
            target_fps: Desired sampling FPS
        """
        self.source_fps = source_fps
        self.target_fps = target_fps
        self.frame_interval = get_frame_interval(source_fps, target_fps)

        self.frame_idx = 0
        self.next_sample_pos = 0.0
        self.sampled_count = 0

    def reset(self):
        """Reset the sampler state."""
        self.frame_idx = 0
        self.next_sample_pos = 0.0
        self.sampled_count = 0

    def should_sample(self) -> Tuple[bool, int]:
        """
        Check if the current frame should be sampled.

        Returns:
            Tuple of (should_sample, sampled_index)
            - should_sample: True if this frame should be sampled
            - sampled_index: The index in the sampled sequence (only valid if should_sample)
        """
        if self.frame_idx >= int(self.next_sample_pos):
            return True, self.sampled_count
        return False, -1

    def advance(self, was_sampled: bool = None):
        """
        Advance to the next frame.

        Args:
            was_sampled: If provided, explicitly indicate if frame was sampled.
                        If None, automatically determined from should_sample state.
        """
        if was_sampled is None:
            was_sampled = self.frame_idx >= int(self.next_sample_pos)

        if was_sampled:
            self.next_sample_pos += self.frame_interval
            self.sampled_count += 1

        self.frame_idx += 1

    def get_timestamp(self) -> float:
        """Get timestamp for current sampled frame in seconds."""
        return calculate_timestamp(self.sampled_count, self.target_fps)

    @property
    def stats(self) -> Dict[str, Any]:
        """Get current sampler statistics."""
        return {
            'frames_read': self.frame_idx,
            'frames_sampled': self.sampled_count,
            'frame_interval': self.frame_interval,
            'next_sample_pos': self.next_sample_pos
        }


# =============================================================================
# DECORD-BASED FUNCTIONS (with OpenCV fallback)
# =============================================================================

def iter_sampled_frames_decord(
    video_path: str,
    target_fps: float,
    resize: Optional[Tuple[int, int]] = None,
    convert_rgb: bool = True,
    use_gpu: bool = False
) -> Generator[Tuple[int, int, np.ndarray], None, None]:
    """
    Generator that yields sampled frames using decord (or OpenCV fallback).

    Uses batch frame reading for efficient video decoding when decord is available.
    Falls back to OpenCV sequential reading otherwise.

    Args:
        video_path: Path to the video file
        target_fps: Desired FPS for sampling
        resize: Optional (width, height) to resize frames
        convert_rgb: Whether to convert to RGB (decord returns RGB, OpenCV returns BGR)
        use_gpu: Use GPU for video decoding (requires CUDA-enabled decord)

    Yields:
        Tuple of (original_frame_idx, sampled_frame_idx, frame)
    """
    if DECORD_AVAILABLE:
        # Use decord for fast frame reading
        if use_gpu:
            try:
                ctx = gpu(0)
            except Exception:
                ctx = cpu(0)
        else:
            ctx = cpu(0)

        vr = VideoReader(video_path, ctx=ctx)
        source_fps = vr.get_avg_fps()
        total_frames = len(vr)

        frame_indices = calculate_frame_indices(total_frames, source_fps, target_fps)

        for sampled_idx, orig_idx in enumerate(frame_indices):
            frame = vr[orig_idx].asnumpy()

            if resize:
                frame = cv2.resize(frame, resize)

            # decord returns RGB by default
            if not convert_rgb:
                frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)

            yield (orig_idx, sampled_idx, frame)
    else:
        # Fallback to OpenCV sequential reading
        yield from iter_sampled_frames_sequential(
            video_path, target_fps, resize, convert_rgb
        )


def read_sampled_frames_batch(
    video_path: str,
    target_fps: float,
    resize: Optional[Tuple[int, int]] = None,
    convert_rgb: bool = True,
    use_gpu: bool = False
) -> Tuple[List[np.ndarray], List[int], Dict[str, Any]]:
    """
    Read all sampled frames at once.

    Uses decord batch reading when available (fastest), otherwise falls back
    to OpenCV sequential reading.

    Args:
        video_path: Path to the video file
        target_fps: Desired FPS for sampling
        resize: Optional (width, height) to resize frames
        convert_rgb: Whether frames should be RGB (True) or BGR (False)
        use_gpu: Use GPU for video decoding

    Returns:
        Tuple of:
        - frames: List of numpy arrays (all sampled frames)
        - indices: List of original frame indices
        - info: Video metadata dict
    """
    if DECORD_AVAILABLE:
        frames = None
        source_fps = None
        total_frames = None
        frame_indices = None

        # Try GPU decode first.
        # GPU decode can deadlock on malformed EOF (missing final GOP — common in dashcams).
        # Python signals cannot interrupt blocked C extensions, so we use a daemon thread:
        # the main thread waits 5s then falls back to CPU; the zombie GPU thread is left
        # to die on its own (daemon=True ensures it dies with the process).
        if use_gpu:
            import threading
            _gpu_result = [None]  # [frames] on success, [None] on failure

            def _gpu_decode_worker():
                try:
                    ctx = gpu(0)
                    vr = VideoReader(video_path, ctx=ctx)
                    src_fps = vr.get_avg_fps()
                    n_frames = len(vr)
                    indices = calculate_frame_indices(n_frames, src_fps, target_fps)
                    batch = vr.get_batch(indices).asnumpy()
                    decoded = []
                    for frame in batch:
                        if resize:
                            frame = cv2.resize(frame, resize)
                        if not convert_rgb:
                            frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
                        decoded.append(frame)
                    _gpu_result[0] = decoded
                except Exception:
                    pass  # leave _gpu_result[0] as None

            t = threading.Thread(target=_gpu_decode_worker, daemon=True)
            t.start()
            t.join(timeout=5)
            if t.is_alive():
                frames = None  # GPU deadlocked — zombie thread will die with process
            else:
                frames = _gpu_result[0]

        # CPU decode: use_gpu=False, or GPU attempt failed
        if frames is None:
            ctx = cpu(0)
            vr = VideoReader(video_path, ctx=ctx)
            source_fps = vr.get_avg_fps()
            total_frames = len(vr)
            frame_indices = calculate_frame_indices(total_frames, source_fps, target_fps)
            frames_batch = vr.get_batch(frame_indices).asnumpy()
            frames = []
            for frame in frames_batch:
                if resize:
                    frame = cv2.resize(frame, resize)
                if not convert_rgb:
                    frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
                frames.append(frame)
    else:
        # Fallback to OpenCV
        video_info = get_video_info(video_path)
        source_fps = video_info['fps']
        total_frames = video_info['total_frames']

        frame_indices = calculate_frame_indices(total_frames, source_fps, target_fps)

        frames = []
        cap = cv2.VideoCapture(video_path)
        try:
            for orig_idx in frame_indices:
                cap.set(cv2.CAP_PROP_POS_FRAMES, orig_idx)
                ret, frame = cap.read()
                if not ret:
                    break

                if resize:
                    frame = cv2.resize(frame, resize)
                if convert_rgb:
                    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                frames.append(frame)
        finally:
            cap.release()

    info = {
        'fps': source_fps,
        'total_frames': total_frames,
        'sampled_frames': len(frames),
        'target_fps': target_fps,
        'frame_interval': get_frame_interval(source_fps, target_fps)
    }

    return frames, frame_indices, info


def iter_sliding_windows_decord(
    video_path: str,
    target_fps: float,
    window_size: int = 16,
    stride: int = 4,
    resize: Optional[Tuple[int, int]] = None,
    convert_rgb: bool = True,
    use_gpu: bool = False
) -> Generator[Tuple[List[np.ndarray], int, int, float], None, None]:
    """
    Iterate over sliding windows of sampled frames.

    Uses decord when available for fast batch reading, otherwise falls back
    to OpenCV. This is optimized for model inference where you need windows
    of consecutive frames with a stride between windows.

    Args:
        video_path: Path to the video file
        target_fps: Desired FPS for sampling
        window_size: Number of frames per window (default: 16)
        stride: Number of frames to advance between windows (default: 4)
        resize: Optional (width, height) to resize frames
        convert_rgb: Whether frames should be RGB
        use_gpu: Use GPU for video decoding

    Yields:
        Tuple of:
        - window_frames: List of numpy arrays (window_size frames)
        - window_idx: Index of this window
        - center_sampled_idx: Index of center frame in sampled sequence
        - timestamp: Timestamp of window center in seconds
    """
    # Read all sampled frames at once (uses decord or OpenCV internally)
    frames, indices, info = read_sampled_frames_batch(
        video_path, target_fps, resize, convert_rgb, use_gpu
    )

    num_sampled = len(frames)
    if num_sampled < window_size:
        # Pad with first frame if video is too short
        while len(frames) < window_size:
            frames.insert(0, frames[0])
            indices.insert(0, indices[0])

        yield (
            frames[:window_size],
            0,
            window_size // 2,
            calculate_timestamp(window_size // 2, target_fps)
        )
        return

    # Generate sliding windows
    window_idx = 0
    start = 0

    while start + window_size <= num_sampled:
        window_frames = frames[start:start + window_size]
        center_idx = start + window_size // 2
        timestamp = calculate_timestamp(center_idx, target_fps)

        yield (window_frames, window_idx, center_idx, timestamp)

        window_idx += 1
        start += stride

    # Note: We don't generate a final partial window to match the app's behavior.
    # The app stops when it can't create another full window with the given stride.
