"""Media validation helpers for MiniMax-H3 reference inputs."""

import io
import os
from typing import Any, Dict, Iterable


MB = 1024 * 1024
MIN_MEDIA_DIMENSION = 256
MAX_MEDIA_DIMENSION = 5760
MIN_MEDIA_RATIO = 0.4
MAX_MEDIA_RATIO = 2.5

MAX_IMAGE_BYTES = 30 * MB
MAX_REFERENCE_IMAGES = 9

MAX_VIDEO_BYTES = 50 * MB
MAX_REFERENCE_VIDEOS = 3
MIN_REFERENCE_DURATION = 2.0
MAX_REFERENCE_DURATION = 15.0
MAX_TOTAL_VIDEO_DURATION = 15.0
MIN_VIDEO_FPS = 23.976
MAX_VIDEO_FPS = 60.0

MAX_AUDIO_BYTES = 15 * MB
MAX_REFERENCE_AUDIOS = 3
MAX_TOTAL_AUDIO_DURATION = 15.0


def _validate_dimensions(width: int, height: int, label: str) -> None:
    if not (
        MIN_MEDIA_DIMENSION <= width <= MAX_MEDIA_DIMENSION
        and MIN_MEDIA_DIMENSION <= height <= MAX_MEDIA_DIMENSION
    ):
        raise ValueError(
            f"{label}宽高必须均在 {MIN_MEDIA_DIMENSION}～{MAX_MEDIA_DIMENSION}px，"
            f"当前为 {width}x{height}。"
        )
    ratio = width / height
    if not MIN_MEDIA_RATIO <= ratio <= MAX_MEDIA_RATIO:
        raise ValueError(
            f"{label}宽高比必须在 0.4～2.5，当前为 {ratio:.3f}（{width}:{height}）。"
        )


def validate_image(image, label: str = "图片") -> Dict[str, Any]:
    """Validate the exact PNG bytes sent by the shared uploader."""
    width, height = image.size
    _validate_dimensions(int(width), int(height), label)

    buffer = io.BytesIO()
    image.save(buffer, format="PNG")
    size = buffer.tell()
    if size > MAX_IMAGE_BYTES:
        raise ValueError(
            f"{label}转为 PNG 后不能超过 30MB，当前为 {size / MB:.2f}MB。"
        )
    return {"width": width, "height": height, "size": size}


def _get_video_source(video):
    if hasattr(video, "get_stream_source"):
        source = video.get_stream_source()
    elif isinstance(video, dict):
        source = (
            video.get("video")
            or video.get("path")
            or video.get("file")
            or video.get("filename")
            or video.get("source_path")
        )
    elif isinstance(video, (str, os.PathLike, io.BytesIO)):
        source = video
    else:
        source = None
        for attr in ("source_path", "path", "video", "file", "filename"):
            if hasattr(video, attr):
                source = getattr(video, attr)
                if source:
                    break
    if source is None:
        raise ValueError("无法获取参考视频数据。")
    return source


def _source_for_probe(source):
    if isinstance(source, io.BytesIO):
        data = source.getvalue()
        return io.BytesIO(data), len(data)
    try:
        path = os.fspath(source)
    except TypeError:
        raise ValueError(f"无法读取参考视频来源：{type(source).__name__}") from None
    if not os.path.isfile(path):
        raise ValueError(f"参考视频文件不存在：{path}")
    return path, os.path.getsize(path)


def probe_video(video) -> Dict[str, Any]:
    """Inspect a ComfyUI VIDEO using PyAV without transcoding it."""
    try:
        import av
    except ImportError:
        raise RuntimeError("当前 ComfyUI 环境缺少 PyAV，无法校验 MiniMax H3 参考视频。") from None

    source = _get_video_source(video)
    probe_source, size = _source_for_probe(source)
    try:
        container = av.open(probe_source)
    except Exception as exc:
        raise ValueError(f"无法读取参考视频：{exc}") from None

    try:
        video_streams = list(container.streams.video)
        if not video_streams:
            raise ValueError("参考视频不包含视频轨道。")
        stream = video_streams[0]
        width = int(stream.codec_context.width or 0)
        height = int(stream.codec_context.height or 0)
        video_codec = str(stream.codec_context.name or "").lower()
        rate = stream.average_rate or stream.base_rate or stream.guessed_rate
        fps = float(rate) if rate else 0.0

        duration = None
        if stream.duration is not None and stream.time_base is not None:
            duration = float(stream.duration * stream.time_base)
        elif container.duration is not None:
            duration = float(container.duration / av.time_base)
        duration = float(duration or 0.0)

        audio_codecs = {
            str(audio_stream.codec_context.name or "").lower()
            for audio_stream in container.streams.audio
        }
        format_names = {
            name.strip().lower()
            for name in str(container.format.name or "").split(",")
            if name.strip()
        }
    finally:
        container.close()

    return {
        "size": size,
        "width": width,
        "height": height,
        "duration": duration,
        "fps": fps,
        "video_codec": video_codec,
        "audio_codecs": audio_codecs,
        "format_names": format_names,
    }


def validate_video_info(info: Dict[str, Any], label: str = "参考视频") -> None:
    size = int(info.get("size") or 0)
    if size <= 0:
        raise ValueError(f"{label}文件为空。")
    if size > MAX_VIDEO_BYTES:
        raise ValueError(f"{label}不能超过 50MB，当前为 {size / MB:.2f}MB。")

    format_names = set(info.get("format_names") or ())
    if not format_names.intersection({"mp4", "mov"}):
        raise ValueError(f"{label}容器仅支持 MP4/MOV，当前为 {sorted(format_names) or '未知'}。")

    duration = float(info.get("duration") or 0.0)
    if not MIN_REFERENCE_DURATION <= duration <= MAX_REFERENCE_DURATION:
        raise ValueError(f"{label}时长必须为 2～15 秒，当前为 {duration:.3f} 秒。")

    fps = float(info.get("fps") or 0.0)
    if not MIN_VIDEO_FPS <= fps <= MAX_VIDEO_FPS:
        raise ValueError(f"{label}帧率必须为 23.976～60 FPS，当前为 {fps:.3f} FPS。")

def validate_reference_videos(videos: Iterable[Any]) -> list[Dict[str, Any]]:
    videos = list(videos)
    if len(videos) > MAX_REFERENCE_VIDEOS:
        raise ValueError(f"参考视频最多 {MAX_REFERENCE_VIDEOS} 个，当前为 {len(videos)} 个。")
    infos = []
    for index, video in enumerate(videos, start=1):
        info = probe_video(video)
        validate_video_info(info, f"参考视频{index}")
        infos.append(info)
    total_duration = sum(float(info["duration"]) for info in infos)
    if total_duration > MAX_TOTAL_VIDEO_DURATION:
        raise ValueError(f"参考视频总时长不能超过 15 秒，当前为 {total_duration:.3f} 秒。")
    return infos


def inspect_audio(audio, label: str = "参考音频") -> Dict[str, Any]:
    if not isinstance(audio, dict):
        raise ValueError(f"{label}数据格式无效。")
    waveform = audio.get("waveform")
    sample_rate = int(audio.get("sample_rate") or 0)
    if waveform is None or sample_rate <= 0:
        raise ValueError(f"{label}缺少 waveform 或 sample_rate。")

    shape = tuple(int(v) for v in waveform.shape)
    if not shape:
        raise ValueError(f"{label}波形为空。")
    if len(shape) == 3 and shape[0] != 1:
        raise ValueError(f"{label}仅支持一个音频批次，当前批次为 {shape[0]}。")
    samples = shape[-1]
    if samples <= 0:
        raise ValueError(f"{label}波形为空。")
    duration = samples / sample_rate

    # The shared uploader currently downmixes to mono 16-bit PCM WAV.
    encoded_size = 44 + samples * 2
    if encoded_size > MAX_AUDIO_BYTES:
        raise ValueError(
            f"{label}编码为 WAV 后不能超过 15MB，预计为 {encoded_size / MB:.2f}MB。"
        )
    if not MIN_REFERENCE_DURATION <= duration <= MAX_REFERENCE_DURATION:
        raise ValueError(f"{label}时长必须为 2～15 秒，当前为 {duration:.3f} 秒。")
    return {
        "sample_rate": sample_rate,
        "samples": samples,
        "duration": duration,
        "encoded_size": encoded_size,
    }


def validate_reference_audios(audios: Iterable[Any]) -> list[Dict[str, Any]]:
    audios = list(audios)
    if len(audios) > MAX_REFERENCE_AUDIOS:
        raise ValueError(f"参考音频最多 {MAX_REFERENCE_AUDIOS} 个，当前为 {len(audios)} 个。")
    infos = [inspect_audio(audio, f"参考音频{index}") for index, audio in enumerate(audios, 1)]
    total_duration = sum(float(info["duration"]) for info in infos)
    if total_duration > MAX_TOTAL_AUDIO_DURATION:
        raise ValueError(f"参考音频总时长不能超过 15 秒，当前为 {total_duration:.3f} 秒。")
    return infos