"""Video frame chunking, overlap calculation, and overlap blending."""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np


@dataclass
class ChunkSpec:
    chunk_id: int
    start_frame: int       # inclusive, with overlap
    end_frame: int         # exclusive, with overlap
    content_start: int     # inclusive, no overlap (for final assembly)
    content_end: int       # exclusive, no overlap

    @property
    def total_frames(self) -> int:
        return self.end_frame - self.start_frame


@dataclass
class ChunkResult:
    spec: ChunkSpec
    frames: list[np.ndarray]  # full chunk frames including overlap


class VideoChunker:
    def __init__(
        self,
        overlap_frames: int = 10,
        blend_mode: str = "linear",
    ) -> None:
        self.overlap_frames = overlap_frames
        self.blend_mode = blend_mode

    @staticmethod
    def compute_overlap(chunk_size: int, neighbor_length: int = 10) -> int:
        """Compute overlap frames based on model params and chunk size.

        Rules:
        - overlap = max(neighbor_length, 5)
        - if chunk_size <= 20: overlap = chunk_size // 4
        - overlap <= chunk_size * 0.4 (S2-R1-012)
        """
        overlap = max(neighbor_length, 5)

        if chunk_size <= 20:
            overlap = max(chunk_size // 4, 2)

        max_overlap = int(chunk_size * 0.4)
        overlap = min(overlap, max_overlap)

        return max(overlap, 1)

    def create_chunks(
        self,
        total_frames: int,
        chunk_size: int,
    ) -> list[ChunkSpec]:
        """Generate chunk specs covering all frames with overlap."""
        if total_frames <= 0:
            return []

        if chunk_size <= 0:
            chunk_size = total_frames

        if total_frames <= chunk_size:
            return [ChunkSpec(
                chunk_id=0,
                start_frame=0,
                end_frame=total_frames,
                content_start=0,
                content_end=total_frames,
            )]

        overlap = self.overlap_frames
        stride = chunk_size - overlap
        if stride <= 0:
            stride = max(1, chunk_size // 2)

        chunks: list[ChunkSpec] = []
        chunk_id = 0
        pos = 0

        while pos < total_frames:
            start = pos
            end = min(start + chunk_size, total_frames)

            if chunk_id == 0:
                content_start = 0
            else:
                content_start = start + overlap // 2

            if end >= total_frames:
                content_end = total_frames
            else:
                content_end = end - overlap // 2

            # Ensure content region is valid
            content_start = min(content_start, content_end)

            chunks.append(ChunkSpec(
                chunk_id=chunk_id,
                start_frame=start,
                end_frame=end,
                content_start=content_start,
                content_end=content_end,
            ))

            if end >= total_frames:
                break

            pos += stride
            chunk_id += 1

        return chunks

    def blend_overlap_pair(
        self,
        prev_tail: list[np.ndarray],
        curr_head: list[np.ndarray],
    ) -> list[np.ndarray]:
        """Blend overlapping region between two consecutive chunks (vectorized)."""
        n = min(len(prev_tail), len(curr_head))
        if n == 0:
            return []

        # Stack into (N, H, W, 3) arrays for vectorized ops
        prev_arr = np.stack(prev_tail[:n]).astype(np.float32)
        curr_arr = np.stack(curr_head[:n]).astype(np.float32)

        # Alpha ramp: shape (N, 1, 1, 1) for broadcasting
        alphas = np.arange(1, n + 1, dtype=np.float32) / (n + 1)
        if self.blend_mode == "sigmoid":
            alphas = 1.0 / (1.0 + np.exp(-12 * (alphas - 0.5)))
        alphas = alphas.reshape(n, 1, 1, 1)

        blended = ((1.0 - alphas) * prev_arr + alphas * curr_arr).clip(0, 255).astype(np.uint8)
        return [blended[i] for i in range(n)]

    def extract_overlap_regions(
        self,
        chunk_frames: list[np.ndarray],
        spec: ChunkSpec,
        prev_spec: ChunkSpec | None,
    ) -> tuple[list[np.ndarray], list[np.ndarray], list[np.ndarray]]:
        """Split chunk frames into (head_overlap, content, tail_overlap).

        head_overlap: frames that overlap with previous chunk
        content: non-overlapping frames
        tail_overlap: frames that overlap with next chunk
        """
        total = len(chunk_frames)

        if prev_spec is None:
            head_size = 0
        else:
            head_size = spec.content_start - spec.start_frame

        tail_size = spec.end_frame - spec.content_end

        head_overlap = chunk_frames[:head_size] if head_size > 0 else []
        tail_overlap = chunk_frames[total - tail_size:] if tail_size > 0 else []
        content = chunk_frames[head_size:total - tail_size if tail_size > 0 else total]

        return head_overlap, content, tail_overlap
