"""Preprocessing helpers for image/video prediction inputs."""

from __future__ import annotations

from dataclasses import dataclass
from io import BytesIO

import cv2
import torch
from PIL import Image
from torchvision import transforms

from .io_utils import temporary_file_from_bytes


@dataclass(frozen=True)
class ImagePreprocessConfig:
    """Standard image preprocessing configuration."""

    size: tuple[int, int]
    mean: tuple[float, float, float] = (0.485, 0.456, 0.406)
    std: tuple[float, float, float] = (0.229, 0.224, 0.225)


def load_rgb_image_from_bytes(image_bytes: bytes) -> Image.Image:
    """Decode image bytes into an RGB PIL image.

    Args:
        image_bytes: Raw encoded image bytes.

    Returns:
        Decoded RGB ``PIL.Image`` object.
    """
    if not image_bytes:
        raise ValueError("`image_bytes` cannot be empty.")
    return Image.open(BytesIO(image_bytes)).convert("RGB")


def build_image_transform(config: ImagePreprocessConfig) -> transforms.Compose:
    """Construct torchvision transform pipeline for model inference.

    Args:
        config: Preprocessing configuration containing size and normalization.

    Returns:
        Composed torchvision transform callable.
    """
    return transforms.Compose(
        [
            transforms.Resize(config.size),
            transforms.ToTensor(),
            transforms.Normalize(mean=list(config.mean), std=list(config.std)),
        ]
    )


def get_model_device(model: torch.nn.Module) -> torch.device:
    """Return model parameter device, or CPU for parameter-less models.

    Args:
        model: Torch module whose parameter device should be inspected.

    Returns:
        Device of first parameter, or CPU if module has no parameters.
    """
    try:
        return next(model.parameters()).device
    except StopIteration:
        return torch.device("cpu")


def preprocess_image_from_bytes(
    image_bytes: bytes,
    *,
    config: ImagePreprocessConfig,
    device: torch.device | None = None,
) -> tuple[torch.Tensor, Image.Image]:
    """Decode and transform image bytes into batched tensor ``[1, C, H, W]``.

    Args:
        image_bytes: Raw encoded image bytes.
        config: Preprocessing configuration for resize and normalization.
        device: Optional device to move the output tensor to.

    Returns:
        Tuple of ``(image_tensor, pil_image)``.
    """
    image = load_rgb_image_from_bytes(image_bytes)
    image_tensor = build_image_transform(config)(image).unsqueeze(0)
    if device is not None:
        image_tensor = image_tensor.to(device)
    return image_tensor, image


def probe_video_fps(video_bytes: bytes, *, default_fps: float = 30.0) -> float:
    """Extract FPS from video bytes using OpenCV with temp-file isolation.

    Args:
        video_bytes: Raw encoded video bytes.
        default_fps: Fallback FPS value when probing fails.

    Returns:
        Probed FPS value, or ``default_fps`` when unavailable.
    """
    with temporary_file_from_bytes(video_bytes, suffix=".mp4") as temp_path:
        capture = cv2.VideoCapture(temp_path)
        if not capture.isOpened():
            capture.release()
            return float(default_fps)
        fps = capture.get(cv2.CAP_PROP_FPS)
        capture.release()
        if fps and fps > 0:
            return float(fps)
    return float(default_fps)
