import enum
import logging
import math
import threading
from collections.abc import Generator, Iterator
from fractions import Fraction
from io import BytesIO
from pathlib import Path
from queue import Queue

import av
import numpy as np
import OpenImageIO
import torch
from einops import rearrange
from PIL import Image
from torch._prims_common import DeviceLikeType
from tqdm import tqdm

from ltx_core.hdr import LogC3
from ltx_core.types import Audio, VideoPixelShape
from ltx_pipelines.utils.color_conversion import FrameConverter, PixelFormat, yuv420p_bt709_converter_
from ltx_pipelines.utils.constants import DEFAULT_IMAGE_CRF

logger = logging.getLogger(__name__)


class ResizeMode(enum.Enum):
    """How to fit a conditioning video to the target resolution."""

    CENTER_CROP = "center_crop"
    REFLECT_PAD = "reflect_pad"


def resize_aspect_ratio_preserving(image: torch.Tensor, long_side: int) -> torch.Tensor:
    """
    Resize image preserving aspect ratio (filling target long side).
    Preserves the input dimensions order.
    Args:
        image: Input image tensor with shape (F (optional), H, W, C)
        long_side: Target long side size.
    Returns:
        Tensor with shape (F (optional), H, W, C) F = 1 if input is 3D, otherwise input shape[0]
    """
    height, width = image.shape[-3:2]
    max_side = max(height, width)
    scale = long_side / float(max_side)
    target_height = int(height * scale)
    target_width = int(width * scale)
    resized = resize_and_center_crop(image, target_height, target_width)
    # rearrange and remove batch dimension
    result = rearrange(resized, "b c f h w -> b f h w c")[0]
    # preserve input dimensions
    return result[0] if result.shape[0] == 1 else result


def resize_and_center_crop(tensor: torch.Tensor, height: int, width: int) -> torch.Tensor:
    """
    Resize tensor preserving aspect ratio (filling target), then center crop to exact dimensions.
    Args:
        latent: Input tensor with shape (H, W, C) or (F, H, W, C)
        height: Target height
        width: Target width
    Returns:
        Tensor with shape (1, C, 1, height, width) for 3D input or (1, C, F, height, width) for 4D input
    """
    if tensor.ndim == 3:
        tensor = rearrange(tensor, "h w c -> 1 c h w")
    elif tensor.ndim == 4:
        tensor = rearrange(tensor, "f h w c -> f c h w")
    else:
        raise ValueError(f"Expected input with 3 or 4 dimensions; got shape {tensor.shape}.")

    _, _, src_h, src_w = tensor.shape

    scale = max(height / src_h, width / src_w)
    # Use ceil to avoid floating-point rounding causing new_h/new_w to be
    # slightly smaller than target, which would result in negative crop offsets.
    new_h = math.ceil(src_h * scale)
    new_w = math.ceil(src_w * scale)

    tensor = torch.nn.functional.interpolate(tensor, size=(new_h, new_w), mode="bilinear", align_corners=False)

    crop_top = (new_h - height) // 2
    crop_left = (new_w - width) // 2
    tensor = tensor[:, :, crop_top : crop_top + height, crop_left : crop_left + width]

    tensor = rearrange(tensor, "f c h w -> 1 c f h w")
    return tensor


def normalize_images(images: torch.Tensor, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
    return (images / 127.5 - 1.0).to(device=device, dtype=dtype)


def to_vae_range(x: torch.Tensor) -> torch.Tensor:
    """Map [0, 1] to [-1, 1] (VAE input convention)."""
    return torch.clamp(x, 0.0, 1.0) * 2.0 - 1.0


def from_vae_range(z: torch.Tensor) -> torch.Tensor:
    """Map [-1, 1] (VAE output convention) to [0, 1]."""
    return torch.clamp((z + 1.0) / 2.0, 0.0, 1.0)


def load_image_and_preprocess(
    image_path: str,
    height: int,
    width: int,
    dtype: torch.dtype,
    device: torch.device,
    crf: int = DEFAULT_IMAGE_CRF,
) -> torch.Tensor:
    """
    Loads an image from a path and preprocesses it for conditioning.
    Note: The image is resized to the nearest multiple of 2 for compatibility with video codecs.
    """
    image = decode_image(image_path=image_path)
    image = preprocess(image=image, crf=crf)
    image = torch.tensor(image, dtype=torch.float32, device=device)
    image = resize_and_center_crop(image, height, width)
    image = normalize_images(image, device, dtype)
    return image


def video_preprocess(
    frames: Generator[torch.Tensor],
    height: int,
    width: int,
    dtype: torch.dtype,
    device: torch.device,
) -> torch.Tensor:
    """Preprocesses a video frame generator for conditioning.
    Args:
        frames: Generator of video frames as tensors of shape (1, H, W, C), dtype uint8.
        height: Target height in pixels.
        width: Target width in pixels.
        dtype: Target dtype for the output tensor.
        device: Target device for the output tensor.
    Returns:
        Tensor of shape (1, C, F, height, width) with values in [-1, 1].
    """
    result: torch.Tensor | None = None
    for f in frames:
        frame = resize_and_center_crop(f.to(torch.float32), height, width)
        frame = normalize_images(frame, device, dtype)
        result = frame if result is None else torch.cat([result, frame], dim=2)
    if result is None:
        raise ValueError("video_preprocess received an empty frame generator; no frames were decoded from the source.")
    return result


def align_resolution(
    width: int,
    height: int,
    resize_mode: ResizeMode,
    divisor: int = 64,
) -> tuple[int, int, int, int]:
    """Compute aligned generation dimensions and crop-back size.
    Args:
        width: Source video width (need not be aligned).
        height: Source video height (need not be aligned).
        resize_mode: CENTER_CROP rounds down; REFLECT_PAD rounds up.
        divisor: Alignment divisor (default 64 for two-stage pipelines).
    Returns:
        ``(gen_width, gen_height, crop_width, crop_height)`` where
        ``gen_*`` are multiples of *divisor* and ``crop_*`` are the
        original dimensions to trim back to after decoding.  When no
        cropping is needed ``crop_*`` equals ``gen_*``.
    """
    if resize_mode is ResizeMode.REFLECT_PAD:
        gen_w = ((width + divisor - 1) // divisor) * divisor
        gen_h = ((height + divisor - 1) // divisor) * divisor
    else:
        gen_w = (width // divisor) * divisor
        gen_h = (height // divisor) * divisor

    crop_w = width if gen_w != width else gen_w
    crop_h = height if gen_h != height else gen_h
    return gen_w, gen_h, crop_w, crop_h


def resize_and_reflect_pad(tensor: torch.Tensor, height: int, width: int) -> torch.Tensor:
    """Resize tensor to fit within target, then reflect-pad to exact dimensions.
    Unlike resize_and_center_crop which stretches and crops, this preserves the
    original aspect ratio and pads the shorter dimension with reflected pixels.
    When the target is already >= the source in both dimensions, interpolation
    is skipped entirely to preserve original pixels.
    Args:
        tensor: Input with shape (H, W, C) or (F, H, W, C)
        height: Target height
        width: Target width
    Returns:
        Tensor with shape (1, C, 1, height, width) for 3D or (1, C, F, height, width) for 4D
    """
    if tensor.ndim == 3:
        tensor = rearrange(tensor, "h w c -> 1 c h w")
    elif tensor.ndim == 4:
        tensor = rearrange(tensor, "f h w c -> f c h w")
    else:
        raise ValueError(f"Expected input with 3 or 4 dimensions; got shape {tensor.shape}.")

    _, _, src_h, src_w = tensor.shape

    if height >= src_h and width >= src_w:
        new_h, new_w = src_h, src_w
    else:
        scale = min(height / src_h, width / src_w)
        new_h = round(src_h * scale)
        new_w = round(src_w * scale)
        tensor = torch.nn.functional.interpolate(tensor, size=(new_h, new_w), mode="bilinear", align_corners=False)

    pad_bottom = height - new_h
    pad_right = width - new_w
    if pad_bottom > 0 or pad_right > 0:
        pad_mode = "reflect" if pad_bottom < new_h and pad_right < new_w else "replicate"
        tensor = torch.nn.functional.pad(tensor, (0, pad_right, 0, pad_bottom), mode=pad_mode)

    tensor = rearrange(tensor, "f c h w -> 1 c f h w")
    return tensor


def load_video_conditioning_hdr(
    video_path: str,
    height: int,
    width: int,
    frame_cap: int,
    dtype: torch.dtype,
    device: torch.device,
    hdr_transform: str = "logc3",
    resize_mode: ResizeMode = ResizeMode.CENTER_CROP,
) -> Iterator[torch.Tensor]:
    """Load a video and yield preprocessed frames for HDR IC-LoRA conditioning.
    Decodes through the standard path and applies the LDR compression that
    matches training. Callers are responsible for providing Rec.709 SDR
    input — the HDR IC-LoRA was trained on that color space.
    Args:
        hdr_transform: LDR-compression name (currently only ``logc3``).
        resize_mode: How to fit the video to the target resolution.
    Yields:
        Per-frame tensors of shape ``(1, C, 1, height, width)``.
    """
    if hdr_transform != "logc3":
        raise ValueError(f"Unsupported HDR transform: {hdr_transform}")

    resize_fn = resize_and_reflect_pad if resize_mode is ResizeMode.REFLECT_PAD else resize_and_center_crop

    for f in decode_video_by_frame(path=video_path, frame_cap=frame_cap, device=device):
        frame = resize_fn(f.to(torch.float32), height, width)
        ldr = (frame / 255.0).clamp(0.0, 1.0)
        compressed = LogC3().compress_ldr(ldr)
        yield to_vae_range(compressed).to(device=device, dtype=dtype)


def decode_image(image_path: str) -> np.ndarray:
    image = Image.open(image_path)
    np_array = np.array(image)[..., :3]
    return np_array


def _write_audio(container: av.container.Container, audio_stream: av.audio.AudioStream, audio: Audio) -> None:
    samples = audio.waveform
    if samples.ndim == 1:
        samples = samples[:, None]

    if samples.shape[1] != 2 and samples.shape[0] == 2:
        samples = samples.T

    if samples.shape[1] != 2:
        raise ValueError(f"Expected samples with 2 channels; got shape {samples.shape}.")

    # Convert to int16 packed for ingestion; resampler converts to encoder fmt.
    if samples.dtype != torch.int16:
        samples = torch.clip(samples, -1.0, 1.0)
        samples = (samples * 32767.0).to(torch.int16)

    frame_in = av.AudioFrame.from_ndarray(
        samples.contiguous().reshape(1, -1).cpu().numpy(),
        format="s16",
        layout="stereo",
    )
    frame_in.sample_rate = audio.sampling_rate

    _resample_audio(container, audio_stream, frame_in)


def _prepare_audio_stream(container: av.container.Container, audio_sample_rate: int) -> av.audio.AudioStream:
    """
    Prepare the audio stream for writing.
    """
    audio_stream = container.add_stream("aac", rate=audio_sample_rate)
    audio_stream.codec_context.sample_rate = audio_sample_rate
    audio_stream.codec_context.layout = "stereo"
    audio_stream.codec_context.time_base = Fraction(1, audio_sample_rate)
    return audio_stream


def _resample_audio(
    container: av.container.Container, audio_stream: av.audio.AudioStream, frame_in: av.AudioFrame
) -> None:
    cc = audio_stream.codec_context

    # Use the encoder's format/layout/rate as the *target*
    target_format = cc.format or "fltp"  # AAC → usually fltp
    target_layout = cc.layout or "stereo"
    target_rate = cc.sample_rate or frame_in.sample_rate

    audio_resampler = av.audio.resampler.AudioResampler(
        format=target_format,
        layout=target_layout,
        rate=target_rate,
    )

    audio_next_pts = 0
    for rframe in audio_resampler.resample(frame_in):
        if rframe.pts is None:
            rframe.pts = audio_next_pts
        audio_next_pts += rframe.samples
        rframe.sample_rate = frame_in.sample_rate
        container.mux(audio_stream.encode(rframe))

    # flush audio encoder
    for packet in audio_stream.encode():
        container.mux(packet)


def encode_video(
    video: torch.Tensor | Iterator[torch.Tensor],
    fps: int,
    audio: Audio | None,
    output_path: str,
    video_chunks_number: int,
    frame_converter: FrameConverter = yuv420p_bt709_converter_,
    crf: int = 19,
    preset: str = "veryfast",
    thread_count: int = 0,
) -> None:
    if isinstance(video, torch.Tensor):
        video = iter([video])

    def convert(chunk: torch.Tensor) -> torch.Tensor:
        return frame_converter(chunk.movedim(-1, -3))

    first_chunk = convert(next(video))

    if frame_converter.pixel_format == PixelFormat.RGB24:
        height, width = first_chunk.shape[-3], first_chunk.shape[-2]
    else:
        height = first_chunk.shape[-2] * 2 // 3
        width = first_chunk.shape[-1]

    container = av.open(output_path, mode="w")
    success = False
    try:
        stream = container.add_stream("libx264", rate=int(fps), options={"crf": str(crf), "preset": preset})
        stream.width = width
        stream.height = height
        stream.pix_fmt = "yuv420p"
        stream.codec_context.thread_count = thread_count
        stream.codec_context.thread_type = "FRAME"
        if frame_converter.color_space is not None:
            stream.codec_context.colorspace = frame_converter.color_space.av_colorspace
        if frame_converter.color_range is not None:
            stream.codec_context.color_range = frame_converter.color_range.av_color_range

        if audio is not None:
            audio_stream = _prepare_audio_stream(container, audio.sampling_rate)

        av_format = frame_converter.pixel_format.av_format

        def cpu_chunks() -> Generator[np.ndarray, None, None]:
            yield first_chunk.to("cpu").numpy()
            for chunk in video:
                yield convert(chunk).to("cpu").numpy()

        _encode_chunks_threaded(
            container=container,
            stream=stream,
            av_format=av_format,
            chunks=cpu_chunks(),
            progress_total=video_chunks_number,
        )

        if audio is not None:
            _write_audio(container, audio_stream, audio)
        success = True
    finally:
        container.close()
        if not success:
            Path(output_path).unlink(missing_ok=True)
    logger.info(f"Video saved to {output_path}")


def encode_audio(audio: Audio, output_path: str) -> None:
    """Save an audio waveform as a 16-bit PCM ``.wav`` file at the source sampling rate.
    Reuses :func:`_write_audio` (the same muxing path used by :func:`encode_video`);
    the only difference is a PCM (``pcm_s16le``) stream in a WAV container instead of
    the AAC stream used for muxed video.
    """
    container = av.open(output_path, mode="w")
    audio_stream = container.add_stream("pcm_s16le", rate=audio.sampling_rate)
    audio_stream.codec_context.sample_rate = audio.sampling_rate
    audio_stream.codec_context.layout = "stereo"
    audio_stream.codec_context.time_base = Fraction(1, audio.sampling_rate)
    try:
        _write_audio(container, audio_stream, audio)
    finally:
        container.close()
    logger.info(f"Audio saved to {output_path}")


def _encode_chunks_threaded(
    container: av.container.Container,
    stream: av.video.stream.VideoStream,
    av_format: str,
    chunks: Iterator[np.ndarray],
    progress_total: int,
) -> None:
    """Run libx264 frame.encode + container.mux on a background thread while
    the caller produces numpy chunks on the current thread. The 1-slot queue
    lets the producer get one chunk ahead (so the next VAE/gather chunk
    overlaps with libx264 encoding the previous chunk) without buffering more
    than one chunk in CPU memory.
    """
    chunk_queue: Queue[np.ndarray | None] = Queue(maxsize=1)
    encoder_error: list[BaseException] = []

    def encoder_worker() -> None:
        error: BaseException | None = None
        while True:
            arr = chunk_queue.get()
            if arr is None:
                break
            if error is not None:
                continue
            try:
                for frame_array in arr:
                    frame = av.VideoFrame.from_ndarray(frame_array, format=av_format)
                    for packet in stream.encode(frame):
                        container.mux(packet)
            except Exception as e:
                error = e
        if error is None:
            try:
                for packet in stream.encode():
                    container.mux(packet)
            except Exception as e:
                error = e
        if error is not None:
            encoder_error.append(error)

    encoder_thread = threading.Thread(target=encoder_worker, name="h264-encoder")
    encoder_thread.start()
    try:
        for arr in tqdm(chunks, total=progress_total):
            chunk_queue.put(arr)
    finally:
        chunk_queue.put(None)
        encoder_thread.join()

    if encoder_error:
        raise encoder_error[0]


_INT_FORMAT_MAX: dict[str, float] = {
    "u8": 128.0,
    "u8p": 128.0,
    "s16": 32768.0,
    "s16p": 32768.0,
    "s32": 2147483648.0,
    "s32p": 2147483648.0,
}


def _audio_frame_to_float(frame: av.AudioFrame) -> np.ndarray:
    """Convert an audio frame to a float32 ndarray with values in [-1, 1] and shape (channels, samples)."""
    fmt = frame.format.name
    arr = frame.to_ndarray().astype(np.float32)
    if fmt in _INT_FORMAT_MAX:
        arr = arr / _INT_FORMAT_MAX[fmt]
    if not frame.format.is_planar:
        # Interleaved formats have shape (1, samples * channels) — reshape to (channels, samples).
        channels = len(frame.layout.channels)
        arr = arr.reshape(-1, channels).T
    return arr


def get_videostream_fps(path: str) -> float:
    """Read video stream FPS."""
    container = av.open(path)
    try:
        video_stream = next(s for s in container.streams if s.type == "video")
        return float(video_stream.average_rate)
    finally:
        container.close()


def get_videostream_metadata(path: str) -> VideoPixelShape:
    """Read video stream metadata as a VideoPixelShape with batch=1.
    If frame count is missing in the container, decodes the stream to count frames.
    Args:
        path: Path to the video file.
    Returns:
        VideoPixelShape with batch=1, frames, height, width, and fps populated from the stream.
    """
    container = av.open(path)
    try:
        video_stream = next(s for s in container.streams if s.type == "video")
        fps = float(video_stream.average_rate)
        num_frames = video_stream.frames or 0
        if num_frames == 0:
            num_frames = sum(1 for _ in container.decode(video_stream))
        width = video_stream.codec_context.width
        height = video_stream.codec_context.height
        return VideoPixelShape(batch=1, frames=num_frames, height=height, width=width, fps=fps)
    finally:
        container.close()


def decode_audio_from_file(
    path: str, device: torch.device, start_time: float = 0.0, max_duration: float | None = None
) -> Audio | None:
    """Decodes audio from a file, optionally seeking to a start time and limiting duration.
    Args:
        path: Path to the audio/video file containing an audio stream.
        device: Device to place the resulting tensor on.
        start_time: Start time in seconds to begin reading audio from.
        max_duration: Maximum audio duration in seconds. If None, reads to end of stream.
    Returns:
        An Audio object with waveform of shape (1, channels, samples), or None if no audio stream.
    """
    container = av.open(path)
    try:
        audio_stream = next(s for s in container.streams if s.type == "audio")
    except StopIteration:
        container.close()
        return None

    sample_rate = audio_stream.rate
    start_pts = int(start_time / audio_stream.time_base)
    end_time = start_time + max_duration if max_duration else audio_stream.duration * audio_stream.time_base
    container.seek(start_pts, stream=audio_stream)

    samples = []
    first_frame_time = None
    for frame in container.decode(audio=0):
        if frame.pts is None:
            continue
        frame_time = float(frame.pts * audio_stream.time_base)
        frame_end = frame_time + frame.samples / frame.sample_rate
        if frame_end < start_time:
            continue
        if frame_time > end_time:
            break
        if first_frame_time is None:
            first_frame_time = frame_time
        samples.append(_audio_frame_to_float(frame))

    container.close()

    if not samples:
        return None

    audio = np.concatenate(samples, axis=-1)

    # Trim samples that fall outside the requested [start_time, start_time + max_duration] window.
    # Audio codecs decode in fixed-size frames whose boundaries may not align with the requested
    # time range, so the first frame can start before start_time and the last frame can end after
    # start_time + max_duration.
    skip_samples = round((start_time - first_frame_time) * sample_rate)
    if skip_samples > 0:
        audio = audio[..., skip_samples:]

    if max_duration is not None:
        max_samples = round(max_duration * sample_rate)
        audio = audio[..., :max_samples]

    waveform = torch.from_numpy(audio).to(device).unsqueeze(0)

    return Audio(waveform=waveform, sampling_rate=sample_rate)


def decode_video_by_frame(
    path: str,
    device: DeviceLikeType,
    starting_frame: int = 0,
    frame_cap: int | None = None,
) -> Generator[torch.Tensor]:
    """Decodes video from a file by sequential frame index, without relying on pts.
    Args:
        path: Path to the video file.
        device: Device to place the resulting tensors on.
        starting_frame: Number of leading frames to skip (default 0).
        frame_cap: Maximum number of frames to yield. If None, no frame limit (default None).
    Yields:
        Frames as tensors of shape (1, H, W, C), dtype uint8.
    """
    container = av.open(path)
    try:
        video_stream = next(s for s in container.streams if s.type == "video")
        for index, frame in enumerate(container.decode(video_stream)):
            if index < starting_frame:
                continue
            tensor = torch.tensor(frame.to_rgb().to_ndarray(), dtype=torch.uint8, device=device).unsqueeze(0)
            yield tensor
            if frame_cap is not None:
                frame_cap -= 1
                if frame_cap == 0:
                    break
    finally:
        container.close()


def decode_video_from_file(
    path: str,
    device: DeviceLikeType,
    start_time: float = 0.0,
    max_duration: float | None = None,
) -> Generator[torch.Tensor]:
    """Decodes video from a file using presentation timestamps for time-based trimming.
    If a frame with no pts is encountered, falls back to :func:`decode_video_by_frame`
    using FPS-derived frame indices.
    Args:
        path: Path to the video file.
        device: Device to place the resulting tensors on.
        start_time: Start time in seconds (default 0.0).
        max_duration: Maximum duration in seconds to decode. If None, reads to end of
            stream (default None).
    Yields:
        Frames as tensors of shape (1, H, W, C), dtype uint8.
    """
    container = av.open(path)
    try:
        video_stream = next(s for s in container.streams if s.type == "video")
        time_base = float(video_stream.time_base)

        if start_time > 0:
            container.seek(int(start_time / time_base), stream=video_stream)

        end_time = start_time + max_duration if max_duration is not None else None

        for frame in container.decode(video_stream):
            # PyAV may leave pts unset when the demuxer does not expose per-frame
            # timestamps (e.g. some raw/elementary streams, stripped or missing
            # metadata, or certain remux paths). Without pts we cannot map frames to
            # wall-clock time, so we fall back to sequential frame indices using the
            # stream's average frame rate.
            if frame.pts is None:
                fps = float(video_stream.average_rate)
                starting_frame = round(start_time * fps)
                frame_cap = round(max_duration * fps) if max_duration is not None else None
                yield from decode_video_by_frame(
                    path=path, device=device, starting_frame=starting_frame, frame_cap=frame_cap
                )
                return
            frame_time = frame.pts * time_base
            if frame_time < start_time:
                continue
            if end_time is not None and frame_time >= end_time:
                break
            yield torch.tensor(frame.to_rgb().to_ndarray(), dtype=torch.uint8, device=device).unsqueeze(0)
    finally:
        container.close()


def encode_single_frame(output_file: str, image_array: np.ndarray, crf: float) -> None:
    container = av.open(output_file, "w", format="mp4")
    try:
        stream = container.add_stream("libx264", rate=1, options={"crf": str(crf), "preset": "veryfast"})
        # Round to nearest multiple of 2 for compatibility with video codecs
        height = image_array.shape[0] // 2 * 2
        width = image_array.shape[1] // 2 * 2
        image_array = image_array[:height, :width]
        stream.height = height
        stream.width = width
        av_frame = av.VideoFrame.from_ndarray(image_array, format="rgb24").reformat(format="yuv420p")
        container.mux(stream.encode(av_frame))
        container.mux(stream.encode())
    finally:
        container.close()


def decode_single_frame(video_file: str) -> np.array:
    container = av.open(video_file)
    try:
        stream = next(s for s in container.streams if s.type == "video")
        frame = next(container.decode(stream))
    finally:
        container.close()
    return frame.to_ndarray(format="rgb24")


def preprocess(image: np.array, crf: float = DEFAULT_IMAGE_CRF) -> np.array:
    if crf == 0:
        return image

    with BytesIO() as output_file:
        encode_single_frame(output_file, image, crf)
        video_bytes = output_file.getvalue()
    with BytesIO(video_bytes) as video_file:
        image_array = decode_single_frame(video_file)
    return image_array


def save_exr_tensor(tensor: torch.Tensor, file_path: str | Path, half: bool = False) -> None:
    """Save a single tensor frame as EXR with linear sRGB colorspace metadata.
    Args:
        tensor: ``[H, W, C]`` or ``[C, H, W]`` float tensor.
        file_path: Output path (e.g. ``frame_0000.exr``).
        half: Force float16 output with ZIP compression.
    """
    if tensor.dim() == 3 and tensor.shape[0] == 3:
        tensor = tensor.permute(1, 2, 0)
    use_half = half or tensor.dtype in (torch.float16, torch.half)
    img_np = np.ascontiguousarray(tensor.cpu().numpy().astype(np.float32))
    file_path = str(file_path)

    h, w = img_np.shape[:2]
    fmt = OpenImageIO.HALF if use_half else OpenImageIO.FLOAT
    spec = OpenImageIO.ImageSpec(w, h, 3, fmt)
    spec.channelnames = ("R", "G", "B")
    spec.attribute("compression", "zip")
    spec.attribute("chromaticities", "float[8]", (0.64, 0.33, 0.30, 0.60, 0.15, 0.06, 0.3127, 0.3290))
    spec.attribute("colorSpace", "sRGB")

    out = OpenImageIO.ImageOutput.create(file_path)
    if out is None:
        raise RuntimeError(
            f"Failed to create EXR writer for '{file_path}'. Ensure OpenImageIO is built with OpenEXR support."
        )
    try:
        if not out.open(file_path, spec):
            raise RuntimeError(f"Failed to open EXR file '{file_path}': {out.geterror()}")
        if not out.write_image(img_np):
            raise RuntimeError(f"Failed to write EXR image '{file_path}': {out.geterror()}")
    finally:
        out.close()


def _linear_to_srgb(x: np.ndarray) -> np.ndarray:
    """Linear -> sRGB OETF per IEC 61966-2-1. Input assumed in [0, 1]."""
    x = np.clip(x, 0.0, 1.0)
    return np.where(x <= 0.0031308, x * 12.92, 1.055 * np.power(x, 1.0 / 2.4) - 0.055)


def encode_exr_sequence_to_mp4(exr_dir: Path, output_mp4: Path, frame_rate: float) -> None:
    """Convert a linear EXR frame sequence to sRGB and encode to H.264 .mp4 via PyAV.
    Exposure is fixed at EV=0 (no gain). Each EXR frame is clamped to [0, 1],
    passed through the sRGB OETF, quantised to 8-bit BGR, and fed to a libx264
    stream (crf 18, yuv420p). ``frame_rate`` is the original source video's
    frame rate so playback matches the input timing.
    """
    import os  # noqa: PLC0415

    os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
    import cv2  # noqa: PLC0415

    exr_files = sorted(exr_dir.glob("frame_*.exr"))
    if not exr_files:
        raise FileNotFoundError(f"No EXR frames found in {exr_dir}")

    container = av.open(str(output_mp4), mode="w")
    stream = container.add_stream("libx264", rate=Fraction(frame_rate).limit_denominator(1000))
    stream.pix_fmt = "yuv420p"
    stream.options = {"crf": "18", "movflags": "+faststart"}

    try:
        for i, exr_path in enumerate(exr_files):
            hdr = cv2.imread(str(exr_path), cv2.IMREAD_UNCHANGED).astype(np.float32)
            sdr = _linear_to_srgb(np.maximum(hdr, 0.0))
            bgr8 = (sdr * 255.0 + 0.5).astype(np.uint8)

            if i == 0:
                stream.height = bgr8.shape[0]
                stream.width = bgr8.shape[1]

            frame = av.VideoFrame.from_ndarray(bgr8, format="bgr24")
            for packet in stream.encode(frame):
                container.mux(packet)

        for packet in stream.encode():
            container.mux(packet)
    finally:
        container.close()