import os
from typing import Iterable

import numpy as np
import torch
from PIL import Image


def _tensor_to_uint8(image: torch.Tensor) -> np.ndarray:
    if image.is_floating_point():
        image = image.clamp(0, 1).mul(255)
    image = image.to(dtype=torch.uint8).cpu()
    array = image.numpy()
    if array.ndim == 3 and array.shape[2] >= 3:
        array = array[:, :, :3]
    return array


def save_frames_to_dir(frames: torch.Tensor, output_dir: str, name_pattern: str) -> None:
    os.makedirs(output_dir, exist_ok=True)
    if not isinstance(frames, torch.Tensor) or frames.dim() != 4:
        raise ValueError("frames must be a 4D torch tensor [N, H, W, C].")

    for index, frame in enumerate(frames, start=1):
        filename = name_pattern % index
        path = os.path.join(output_dir, filename)
        array = _tensor_to_uint8(frame)
        image = Image.fromarray(array)
        image.save(path)


def load_frames_from_dir(frames_dir: str, extension: str = ".png") -> torch.Tensor:
    if not os.path.isdir(frames_dir):
        raise ValueError(f"Frames directory does not exist: {frames_dir}")

    files = sorted(
        f for f in os.listdir(frames_dir) if f.lower().endswith(extension)
    )
    if not files:
        raise ValueError("No frames were generated by FFmpeg.")

    images = []
    for file in files:
        path = os.path.join(frames_dir, file)
        with Image.open(path) as image:
            image = image.convert("RGB")
            array = np.asarray(image).astype(np.float32) / 255.0
        images.append(torch.from_numpy(array))

    return torch.stack(images, dim=0)