import random
import re
from pathlib import Path


import numpy as np
from einops import rearrange
import torch
import torchvision.io
from src.config import VAE_TEMPORAL_FACTOR


def seed_everything(seed: int | None = None) -> None:
    if seed is not None:
        random.seed(seed)
        np.random.seed(seed)
        torch.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)


def save_video(tensor: torch.Tensor, save_path: str, fps: int = 30) -> None:
    """Save (T, 3, H, W) pixels as video, including one-frame clips."""
    if tensor.ndim != 4 or tensor.shape[1] != 3 or min(tensor.shape) <= 0:
        raise ValueError(f"Video must have shape (T, 3, H, W), got {tuple(tensor.shape)}.")
    if fps <= 0:
        raise ValueError("Video FPS must be positive.")
    if tensor.is_floating_point():
        tensor = tensor.mul(255).clamp(0, 255).to(torch.uint8)
    elif tensor.dtype != torch.uint8:
        raise ValueError("Video pixels must be uint8 or floating-point values in [0, 1].")
    video_tensor = rearrange(tensor.cpu(), "t c h w -> t h w c")
    torchvision.io.write_video(save_path, video_tensor, fps=fps, options={"crf": "8"})


def get_from_dict(data: dict, path: str):
    """Look up a dot-separated key path."""
    if not isinstance(data, dict):
        raise TypeError(f"Expected a dictionary, got {type(data).__name__}")

    keys = path.split(".")
    current_level = data

    for key in keys:
        current_level = current_level[key]

    return current_level


def find_latest_checkpoint(ckpt_dir: Path) -> Path | None:
    if not ckpt_dir.exists():
        return None

    step_pattern = re.compile(r"global_step(\d+)")
    latest_step = -1
    latest_ckpt = None

    for subdir in ckpt_dir.iterdir():
        if subdir.is_dir():
            match = step_pattern.fullmatch(subdir.name)
            if match:
                step = int(match.group(1))
                checkpoint = subdir / f"step_{step}.pth"
                if step > latest_step and checkpoint.is_file() and checkpoint.stat().st_size > 0:
                    latest_step = step
                    latest_ckpt = subdir

    return latest_ckpt.resolve() if latest_ckpt else None


def _resolve_full_length_target_frames(
    total_frames: int,
    source_fps: float | None,
    target_fps: float | None,
) -> int:
    """Resample the source duration and floor to a valid 8n+1 frame count."""
    if total_frames <= 0:
        raise ValueError("A source video must contain at least one frame.")
    if source_fps and target_fps and source_fps > 0 and target_fps > 0:
        resampled = int(round((total_frames - 1) * target_fps / source_fps)) + 1
    else:
        resampled = total_frames
    return max(1, ((resampled - 1) // VAE_TEMPORAL_FACTOR) * VAE_TEMPORAL_FACTOR + 1)