from __future__ import annotations

from typing import Any

import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image


def _frame_to_tensor(frame: Any) -> torch.Tensor:
    if isinstance(frame, Image.Image):
        arr = np.asarray(frame.convert("RGB"), dtype=np.float32) / 255.0
        return torch.from_numpy(arr).permute(2, 0, 1)
    if isinstance(frame, np.ndarray):
        arr = frame.astype(np.float32, copy=False)
        if arr.max() > 1.0:
            arr = arr / 255.0
        tensor = torch.from_numpy(arr)
        if tensor.ndim == 3 and tensor.shape[-1] == 3:
            tensor = tensor.permute(2, 0, 1)
        if tensor.ndim != 3:
            raise ValueError(f"Expected frame array with 3 dims, got {arr.shape}")
        return tensor.float().contiguous()
    if torch.is_tensor(frame):
        tensor = frame.detach().float()
        if tensor.ndim == 3 and tensor.shape[-1] == 3:
            tensor = tensor.permute(2, 0, 1)
        if tensor.ndim != 3:
            raise ValueError(f"Expected frame tensor with 3 dims, got {tuple(tensor.shape)}")
        if tensor.max() > 1.0:
            tensor = tensor / 255.0
        return tensor.contiguous()
    raise TypeError(f"Unsupported frame type: {type(frame)!r}")


def video_to_t_c_h_w(video: Any) -> torch.Tensor:
    if torch.is_tensor(video):
        tensor = video.detach().float()
        if tensor.ndim == 5 and tensor.shape[0] == 1:
            tensor = tensor.squeeze(0)
        if tensor.ndim == 3:
            tensor = tensor.unsqueeze(0)
        if tensor.ndim != 4:
            raise ValueError(f"Expected video tensor with 4 dims, got {tuple(tensor.shape)}")
        if tensor.shape[-1] == 3:
            tensor = tensor.permute(0, 3, 1, 2)
        if tensor.max() > 1.0:
            tensor = tensor / 255.0
        return tensor.contiguous()

    if isinstance(video, np.ndarray):
        arr = video.astype(np.float32, copy=False)
        if arr.ndim == 5 and arr.shape[0] == 1:
            arr = arr[0]
        if arr.ndim == 3:
            return _frame_to_tensor(arr).unsqueeze(0).contiguous()
        if arr.ndim != 4:
            raise ValueError(f"Expected video array with 4 dims, got {arr.shape}")
        tensor = torch.from_numpy(arr)
        if tensor.shape[-1] == 3:
            tensor = tensor.permute(0, 3, 1, 2)
        if tensor.max() > 1.0:
            tensor = tensor.float() / 255.0
        return tensor.float().contiguous()

    if isinstance(video, (list, tuple)):
        if not video:
            raise ValueError("Video frame list is empty.")
        return torch.stack([_frame_to_tensor(frame) for frame in video], dim=0).contiguous()

    if isinstance(video, Image.Image):
        return _frame_to_tensor(video).unsqueeze(0).contiguous()

    raise TypeError(f"Unsupported video type: {type(video)!r}")


def videos_to_b_t_c_h_w(videos: list[Any], *, size: int | tuple[int, int] | None = None, num_frames: int | None = None) -> torch.Tensor:
    tensors = [video_to_t_c_h_w(video) for video in videos]
    if num_frames is not None:
        tensors = [_match_frames(tensor, num_frames) for tensor in tensors]
    if size is not None:
        target_size = (size, size) if isinstance(size, int) else tuple(size)
        tensors = [_resize_video(tensor, target_size) for tensor in tensors]
    return torch.stack(tensors, dim=0).clamp(0.0, 1.0).contiguous()


def _match_frames(video: torch.Tensor, num_frames: int) -> torch.Tensor:
    if video.shape[0] == num_frames:
        return video
    if video.shape[0] > num_frames:
        idx = torch.linspace(0, video.shape[0] - 1, steps=num_frames).round().long()
        return video[idx]
    reps = int(np.ceil(num_frames / video.shape[0]))
    return video.repeat((reps, 1, 1, 1))[:num_frames]


def _resize_video(video: torch.Tensor, size: tuple[int, int]) -> torch.Tensor:
    return F.interpolate(video, size=size, mode="bilinear", align_corners=False)