"""Shared audio primitives for LTX-2.5 task tabs.

This module owns raw-audio decoding/preprocessing and the narrow Diffusers
adaptation required to treat an audio modality as frozen conditioning. Task
semantics and evidence remain in the individual tab module.
"""
from __future__ import annotations

from contextlib import contextmanager
from enum import StrEnum
from pathlib import Path

import numpy as np
import torch
from diffusers import FlowMatchEulerDiscreteScheduler


_INT_FORMAT_MAX = {
    "u8": 255.0,
    "u8p": 255.0,
    "s16": 32768.0,
    "s16p": 32768.0,
    "s32": 2147483648.0,
    "s32p": 2147483648.0,
}


class AudioAuthority(StrEnum):
    """Which waveform is authoritative for a task's final audio output.

    This is deliberately a small vocabulary, not a task framework. Individual
    task modules continue to own conditioning, diffusion and mux semantics.
    """

    MODEL_GENERATED_JOINT = "model_generated_joint"
    EXTERNAL_FIXED_CONDITIONING_PASSTHROUGH = "external_fixed_conditioning_passthrough"
    SOURCE_VIDEO_PASSTHROUGH = "source_video_passthrough"
    EXTERNAL_REPLACEMENT = "external_replacement"
    GENERATED_AUDIO_OUTPUT = "generated_audio_output"


_AUDIO_AUTHORITY_MEANINGS = {
    AudioAuthority.MODEL_GENERATED_JOINT: "model-generated joint A/V audio is authoritative",
    AudioAuthority.EXTERNAL_FIXED_CONDITIONING_PASSTHROUGH: (
        "external audio conditions generation while its original decoded waveform remains authoritative"
    ),
    AudioAuthority.SOURCE_VIDEO_PASSTHROUGH: "source-video audio is preserved unchanged as final output audio",
    AudioAuthority.EXTERNAL_REPLACEMENT: "an explicit external replacement waveform is authoritative",
    AudioAuthority.GENERATED_AUDIO_OUTPUT: "the task's generated audio waveform is authoritative",
}


def audio_authority_metadata(authority: AudioAuthority | str) -> dict[str, str]:
    """Return a stable diagnostic record for one authority vocabulary value."""
    resolved = AudioAuthority(authority)
    return {"kind": resolved.value, "meaning": _AUDIO_AUTHORITY_MEANINGS[resolved]}


def _audio_frame_to_float(frame) -> np.ndarray:
    """Convert a PyAV AudioFrame to float32 [channels, samples] in [-1, 1]."""
    arr = frame.to_ndarray().astype(np.float32)
    maximum = _INT_FORMAT_MAX.get(frame.format.name)
    if maximum is not None:
        arr = arr / maximum
    if not frame.format.is_planar:
        channels = len(frame.layout.channels)
        arr = arr.reshape(-1, channels).T
    return arr


def decode_audio_file(path: str | Path, *, max_duration: float) -> tuple[np.ndarray, int, dict]:
    """Decode the beginning of an audio/video file using PyAV.

    Returns stereo float32 waveform [2, samples], source sample rate and metadata.
    The first A2V probe intentionally requires enough source audio for the whole
    requested clip instead of silently padding missing content.
    """
    import av

    source_path = Path(path)
    container = av.open(str(source_path))
    try:
        stream = next((item for item in container.streams if item.type == "audio"), None)
        if stream is None:
            raise ValueError("The uploaded file has no audio stream.")
        sample_rate = int(stream.rate)
        samples = []
        decoded_frames = 0
        for frame in container.decode(stream):
            samples.append(_audio_frame_to_float(frame))
            decoded_frames += 1
            if sum(item.shape[-1] for item in samples) >= int(round(float(max_duration) * sample_rate)):
                break
    finally:
        container.close()

    if not samples:
        raise ValueError("The uploaded audio decoded zero samples.")
    waveform = np.concatenate(samples, axis=-1).astype(np.float32, copy=False)
    required_samples = int(round(float(max_duration) * sample_rate))
    if waveform.shape[-1] < required_samples:
        actual_seconds = waveform.shape[-1] / float(sample_rate)
        raise ValueError(
            f"Input audio is too short for this probe: {actual_seconds:.3f}s available, "
            f"{float(max_duration):.3f}s required."
        )
    waveform = waveform[..., :required_samples]
    source_channels = int(waveform.shape[0])
    if source_channels == 1:
        waveform = np.repeat(waveform, 2, axis=0)
        channel_policy = "mono duplicated to stereo"
    elif source_channels >= 2:
        waveform = waveform[:2]
        channel_policy = "first two channels preserved" if source_channels > 2 else "stereo preserved"
    else:
        raise ValueError(f"Unsupported decoded channel count: {source_channels}")
    waveform = np.clip(waveform, -1.0, 1.0).astype(np.float32, copy=False)
    return waveform, sample_rate, {
        "source_channels": source_channels,
        "condition_channels": 2,
        "channel_policy": channel_policy,
        "decoded_frames": decoded_frames,
        "source_sample_rate": sample_rate,
        "trimmed_samples": int(waveform.shape[-1]),
        "trimmed_duration_seconds": waveform.shape[-1] / float(sample_rate),
    }


def waveform_to_log_mel(
    waveform: np.ndarray,
    source_sample_rate: int,
    *,
    target_sample_rate: int,
    mel_bins: int,
    mel_hop_length: int,
    n_fft: int = 1024,
) -> tuple[np.ndarray, dict]:
    """Match the official LTX AudioProcessor frontend on CPU.

    Uses TorchAudio only for resampling/MelSpectrogram; no TorchCodec I/O path is
    used because input decoding is handled by PyAV.
    """
    import torchaudio

    tensor = torch.from_numpy(np.asarray(waveform, dtype=np.float32)).unsqueeze(0)
    if int(source_sample_rate) != int(target_sample_rate):
        tensor = torchaudio.functional.resample(tensor, int(source_sample_rate), int(target_sample_rate))
    transform = torchaudio.transforms.MelSpectrogram(
        sample_rate=int(target_sample_rate),
        n_fft=int(n_fft),
        win_length=int(n_fft),
        hop_length=int(mel_hop_length),
        f_min=0.0,
        f_max=float(target_sample_rate) / 2.0,
        n_mels=int(mel_bins),
        window_fn=torch.hann_window,
        center=True,
        pad_mode="reflect",
        power=1.0,
        mel_scale="slaney",
        norm="slaney",
    )
    mel = transform(tensor)
    mel = torch.log(torch.clamp(mel, min=1e-5))
    # Official ltx-core AudioProcessor returns [batch, channels, time, n_mels].
    mel = mel.permute(0, 1, 3, 2).contiguous().float().cpu()
    return mel.numpy(), {
        "frontend": "official-compatible log-mel",
        "target_sample_rate": int(target_sample_rate),
        "mel_bins": int(mel_bins),
        "mel_hop_length": int(mel_hop_length),
        "n_fft": int(n_fft),
        "mel_shape": [int(value) for value in mel.shape],
    }


def expected_audio_latent_frames(*, num_frames: int, frame_rate: float, sample_rate: int, hop_length: int, compression: int) -> int:
    duration_seconds = float(num_frames) / float(frame_rate)
    per_second = float(sample_rate) / float(hop_length) / float(compression)
    return int(round(duration_seconds * per_second))


def encode_clean_audio_latents(pipe, mel: np.ndarray, *, expected_frames: int) -> torch.Tensor:
    """Encode precomputed log-mel input with Diffusers' already-resident audio VAE."""
    mel_tensor = torch.from_numpy(np.asarray(mel, dtype=np.float32)).to(
        device=pipe._execution_device, dtype=pipe.audio_vae.dtype
    )
    with torch.inference_mode():
        posterior = pipe.audio_vae.encode(mel_tensor).latent_dist
        latents = posterior.mode()
    if latents.shape[2] < int(expected_frames):
        raise RuntimeError(
            f"Encoded audio latent is shorter than the requested video: {latents.shape[2]} < {int(expected_frames)} frames."
        )
    return latents[:, :, : int(expected_frames), :].contiguous()


class _FrozenAudioScheduler(FlowMatchEulerDiscreteScheduler):
    """Audio-only deepcopy of the video scheduler whose step leaves the sample unchanged."""

    def step(self, model_output, timestep, sample, *args, return_dict: bool = True, **kwargs):
        if return_dict:
            # The LTX2 pipeline always requests return_dict=False. Keep failure explicit
            # rather than constructing a scheduler-specific output type here.
            raise RuntimeError("Frozen audio scheduler expects return_dict=False.")
        return (sample,)


class FrozenAudioCopyScheduler(FlowMatchEulerDiscreteScheduler):
    """Normal video scheduler that deep-copies into a frozen audio scheduler.

    LTX2Pipeline currently creates its audio scheduler with copy.deepcopy(self.scheduler).
    This preserves normal video scheduler behaviour while making only that private audio
    copy immutable for A2V.
    """

    def __deepcopy__(self, memo):
        frozen = _FrozenAudioScheduler.from_config(self.config)
        memo[id(self)] = frozen
        return frozen


def frozen_audio_scheduler_from_config(config, **kwargs) -> FrozenAudioCopyScheduler:
    return FrozenAudioCopyScheduler.from_config(config, **kwargs)


@contextmanager
def frozen_audio_conditioning(pipe, *, scheduler: FrozenAudioCopyScheduler):
    """Give one Diffusers LTX2 call the native A2V frozen-audio contract.

    Diffusers exposes the right multimodal model pieces, but its ordinary I2V
    pipeline denoises both modalities. Native A2V instead keeps the input audio
    latent fixed while video is denoised. This task-scoped adapter therefore:

    - normalizes/packs supplied audio latents without adding noise;
    - forces audio timestep/sigma to zero;
    - keeps video scale/shift on the video timestep while overriding only the
      A2V/V2A cross-attention gate inputs to the opposite modality sigma;
    - replaces only the pipeline's private audio scheduler copy with an
      identity-step scheduler.

    All hooks and monkeypatches are removed in ``finally``. The yielded mutable
    evidence dict is filled with hook-count/scale metadata and can be recorded
    by the task module after the context exits.
    """
    original_scheduler = pipe.scheduler
    original_prepare_audio_latents = pipe.prepare_audio_latents
    transformer = pipe.transformer
    video_gate = getattr(transformer, "av_cross_attn_video_a2v_gate", None)
    audio_gate = getattr(transformer, "av_cross_attn_audio_v2a_gate", None)
    if video_gate is None or audio_gate is None:
        raise RuntimeError("Current LTX2 transformer does not expose the expected A2V/V2A gate modules.")

    config = transformer.config
    timestep_multiplier = float(getattr(config, "timestep_scale_multiplier", 1.0) or 1.0)
    cross_multiplier = float(getattr(config, "cross_attn_timestep_scale_multiplier", timestep_multiplier))
    gate_scale_ratio = cross_multiplier / timestep_multiplier
    state = {"video_sigma": None}
    audio_self_attn_modules = [
        block.audio_attn1
        for block in getattr(transformer, "transformer_blocks", [])
        if getattr(block, "audio_attn1", None) is not None
    ]

    def _hook_modules():
        return [transformer, video_gate, audio_gate, *audio_self_attn_modules]

    def _hook_count() -> int:
        return sum(len(getattr(module, "_forward_pre_hooks", {})) for module in _hook_modules())

    baseline_hook_count = _hook_count()
    evidence = {
        "baseline_forward_pre_hook_count": baseline_hook_count,
        "cross_attn_gate_scale_ratio": gate_scale_ratio,
        "video_gate_input": "frozen audio sigma = 0",
        "audio_gate_input": "current video sigma",
        "scale_shift_policy": "per-modality timestep; use_cross_timestep=False",
        "cfg_audio_context_policy": "positive audio context reused for the video-CFG unconditioned half",
        "stg_audio_policy": "audio self-attention remains unperturbed; STG perturbs video only",
        "cfg_audio_context_replacements": 0,
        "audio_self_attention_unperturb_hooks": len(audio_self_attn_modules),
    }

    def _prepare_audio_latents_without_noise(*args, **kwargs):
        kwargs["noise_scale"] = 0.0
        return original_prepare_audio_latents(*args, **kwargs)

    def _freeze_audio_timestep(module, args, kwargs):
        video_timestep = kwargs.get("timestep")
        video_sigma = kwargs.get("sigma")
        audio_timestep = kwargs.get("audio_timestep")
        audio_sigma = kwargs.get("audio_sigma")
        state["video_sigma"] = video_sigma if video_sigma is not None else video_timestep

        # Diffusers I2V uses a per-token video timestep (B, T_video), while
        # audio_timestep/audio_sigma are per-batch (B,). Frozen A2V must zero
        # the existing audio tensors without inheriting the video-token shape.
        # Native LTX keeps Modality.sigma per-batch and derives per-token
        # timesteps separately from the modality denoise mask.
        if isinstance(audio_timestep, torch.Tensor):
            kwargs["audio_timestep"] = torch.zeros_like(audio_timestep)
        elif isinstance(video_sigma, torch.Tensor):
            kwargs["audio_timestep"] = torch.zeros_like(video_sigma)
        elif isinstance(video_timestep, torch.Tensor):
            fallback = video_timestep
            if fallback.ndim > 1:
                fallback = fallback.reshape(fallback.shape[0], -1)[:, 0]
            kwargs["audio_timestep"] = torch.zeros_like(fallback)

        if isinstance(audio_sigma, torch.Tensor):
            kwargs["audio_sigma"] = torch.zeros_like(audio_sigma)
        elif isinstance(video_sigma, torch.Tensor):
            kwargs["audio_sigma"] = torch.zeros_like(video_sigma)
        elif isinstance(kwargs.get("audio_timestep"), torch.Tensor):
            kwargs["audio_sigma"] = torch.zeros_like(kwargs["audio_timestep"])

        # Native A2V uses a positive-only audio guider. When Diffusers batches
        # [video-uncond, video-cond] for CFG, reuse the positive audio prompt
        # context in both halves so only the video text context is unconditioned.
        audio_context = kwargs.get("audio_encoder_hidden_states")
        if (
            isinstance(audio_context, torch.Tensor)
            and audio_context.ndim >= 1
            and audio_context.shape[0] == 2
            and isinstance(video_timestep, torch.Tensor)
            and video_timestep.shape[0] == 2
        ):
            positive_audio_context = audio_context[1:2]
            kwargs["audio_encoder_hidden_states"] = torch.cat(
                [positive_audio_context, positive_audio_context], dim=0
            )
            audio_mask = kwargs.get("audio_encoder_attention_mask")
            if isinstance(audio_mask, torch.Tensor) and audio_mask.shape[0] == 2:
                positive_audio_mask = audio_mask[1:2]
                kwargs["audio_encoder_attention_mask"] = torch.cat(
                    [positive_audio_mask, positive_audio_mask], dim=0
                )
            evidence["cfg_audio_context_replacements"] += 1
        return args, kwargs

    def _replace_first_arg(args, replacement):
        if not args:
            raise RuntimeError("LTX2 cross-attention gate hook received no positional timestep input.")
        return (replacement,) + tuple(args[1:])

    def _zero_video_gate(module, args):
        reference = args[0]
        return _replace_first_arg(args, torch.zeros_like(reference))

    def _video_sigma_for_audio_gate(module, args):
        reference = args[0]
        sigma = state.get("video_sigma")
        if sigma is None:
            raise RuntimeError("Video sigma was not captured before the V2A gate was evaluated.")
        sigma = sigma.to(device=reference.device, dtype=reference.dtype) * gate_scale_ratio
        if sigma.shape != reference.shape:
            try:
                sigma = torch.broadcast_to(sigma, reference.shape)
            except RuntimeError as exc:
                raise RuntimeError(
                    f"Cannot broadcast video sigma shape {tuple(sigma.shape)} to V2A gate input {tuple(reference.shape)}."
                ) from exc
        return _replace_first_arg(args, sigma)

    def _keep_audio_self_attention_unperturbed(module, args, kwargs):
        if "perturbation_mask" in kwargs:
            kwargs["perturbation_mask"] = None
            kwargs["all_perturbed"] = False
        return args, kwargs

    transformer_hook = transformer.register_forward_pre_hook(_freeze_audio_timestep, with_kwargs=True)
    video_gate_hook = video_gate.register_forward_pre_hook(_zero_video_gate)
    audio_gate_hook = audio_gate.register_forward_pre_hook(_video_sigma_for_audio_gate)
    audio_self_attn_hooks = [
        module.register_forward_pre_hook(_keep_audio_self_attention_unperturbed, with_kwargs=True)
        for module in audio_self_attn_modules
    ]
    pipe.scheduler = scheduler
    pipe.prepare_audio_latents = _prepare_audio_latents_without_noise
    evidence["active_forward_pre_hook_count"] = _hook_count()
    try:
        yield evidence
    finally:
        pipe.prepare_audio_latents = original_prepare_audio_latents
        pipe.scheduler = original_scheduler
        for hook in reversed(audio_self_attn_hooks):
            hook.remove()
        audio_gate_hook.remove()
        video_gate_hook.remove()
        transformer_hook.remove()
        restored_hook_count = _hook_count()
        evidence["restored_forward_pre_hook_count"] = restored_hook_count
        evidence["hooks_restored"] = restored_hook_count == baseline_hook_count
        if restored_hook_count != baseline_hook_count:
            raise RuntimeError(
                "Frozen-audio transformer hooks were not fully restored: "
                f"baseline={baseline_hook_count}, restored={restored_hook_count}."
            )