"""프레임 PNG를 디스크에 풀지 않고 메모리에서 디코드·인코드하는 비디오 유틸."""
import io
from collections.abc import Iterable, Iterator, Sequence
from fractions import Fraction
from pathlib import Path

import av
import numpy as np
from PIL import Image

AAC_FRAME_SIZE = 1024


def open_clip(clip: bytes | Path):
    return av.open(io.BytesIO(clip) if isinstance(clip, bytes) else str(clip))


def duration(clip: bytes | Path) -> float:
    container = open_clip(clip)
    try:
        if container.duration is None:
            raise RuntimeError("클립 길이를 읽을 수 없다")
        return container.duration / av.time_base
    finally:
        container.close()


def frames_at(clip: bytes | Path, times: Sequence[float]) -> list[Image.Image]:
    """지정한 시각들의 프레임. seek은 키프레임 단위라 그 이후 첫 프레임을 취한다."""
    container = open_clip(clip)
    stream = container.streams.video[0]
    frames = []
    try:
        for at in times:
            container.seek(int(at / stream.time_base), stream=stream)
            picked = None
            for frame in container.decode(stream):
                picked = frame
                if frame.time is not None and frame.time >= at - 1e-3:
                    break
            if picked is None:
                raise RuntimeError(f"t={at}s 프레임을 못 뽑았다")
            frames.append(picked.to_image())
    finally:
        container.close()
    return frames


def decode_frames(clip: bytes | Path) -> Iterator[Image.Image]:
    """모든 프레임을 순서대로 흘려보낸다. 한 장씩만 메모리에 올라간다."""
    container = open_clip(clip)
    try:
        for frame in container.decode(video=0):
            yield frame.to_image()
    finally:
        container.close()


def encode_mp4(frames: Iterable[Image.Image], size: tuple[int, int],
               audio: np.ndarray | None = None, *, fps: int = 24,
               sample_rate: int = 44100, crf: int = 18,
               audio_bit_rate: int = 192000) -> bytes:
    """H.264 영상 + AAC 스테레오 오디오를 mp4 바이트로 뽑는다. 모노 오디오는 양 채널에 복제한다."""
    buffer = io.BytesIO()
    container = av.open(buffer, "w", format="mp4")

    video_stream = container.add_stream("libx264", rate=fps)
    video_stream.width, video_stream.height = size
    video_stream.pix_fmt = "yuv420p"
    video_stream.options = {"crf": str(crf), "preset": "medium"}

    audio_stream = None
    if audio is not None:
        audio_stream = container.add_stream("aac", rate=sample_rate)
        audio_stream.layout = "stereo"
        audio_stream.bit_rate = audio_bit_rate

    for index, image in enumerate(frames):
        frame = av.VideoFrame.from_image(image)
        frame.pts = index
        for packet in video_stream.encode(frame):
            container.mux(packet)
    for packet in video_stream.encode():
        container.mux(packet)

    if audio_stream is not None:
        interleaved = np.repeat(audio, 2)
        block_size = audio_stream.frame_size or AAC_FRAME_SIZE
        for start in range(0, len(audio), block_size):
            block = interleaved[start * 2:(start + block_size) * 2]
            if len(block) < block_size * 2:
                block = np.pad(block, (0, block_size * 2 - len(block)))
            frame = av.AudioFrame.from_ndarray(block.reshape(1, -1),
                                               format="s16", layout="stereo")
            frame.rate = sample_rate
            frame.pts = start
            frame.time_base = Fraction(1, sample_rate)
            for packet in audio_stream.encode(frame):
                container.mux(packet)
        for packet in audio_stream.encode():
            container.mux(packet)

    container.close()
    return buffer.getvalue()