"""Project-local inference memory optimizations for the S2-Pro DAC codec.

The pinned Fish Speech source remains unchanged. This module builds the same
codec and loads the same checkpoint, then removes buffers that are unnecessary
for the window-limited inference path before the model is moved to CUDA.
"""

from __future__ import annotations

import gc
import io
import math
import os
import threading
import time
from pathlib import Path
from typing import Any

import torch

os.environ.setdefault("TEMPORARILY_DISABLE_PROTOBUF_VERSION_CHECK", "true")

from fish_speech.models.dac.modded_dac import DAC


def load_reference_audio_soundfile(
    reference_audio: bytes | str | Path,
    sample_rate: int,
):
    """Decode API reference audio without TorchCodec.

    Torch 2.11 routes ``torchaudio.load`` through optional TorchCodec. The
    pinned container already includes SoundFile, which supports the WAV/FLAC
    inputs accepted by this service, so no environment package mutation is
    needed.
    """

    import numpy as np
    import soundfile as sf
    import torchaudio

    source = (
        io.BytesIO(reference_audio)
        if isinstance(reference_audio, bytes)
        else reference_audio
    )
    audio, original_rate = sf.read(source, dtype="float32", always_2d=True)
    mono = np.asarray(audio.mean(axis=1), dtype=np.float32)
    if original_rate != sample_rate:
        mono = (
            torchaudio.functional.resample(
                torch.from_numpy(mono),
                original_rate,
                sample_rate,
            )
            .contiguous()
            .numpy()
        )
    return mono


@torch.inference_mode()
def warm_reference_encoder(
    codec: torch.nn.Module,
    device: str | torch.device,
    seconds: float = 1.0,
) -> dict[str, Any]:
    """Prime lazy codec state before the first user reference is cached.

    The staged BF16 encoder produces a different discrete code sequence on its
    first CUDA pass. A discarded silence pass makes subsequent encodes bit
    stable, preventing the first uploaded voice from being conditioned on an
    avoidable cold-start code path.
    """

    if seconds <= 0:
        raise ValueError("Reference warmup duration must be positive")
    target = torch.device(device)
    sample_rate = int(codec.sample_rate)
    samples = int(round(sample_rate * seconds))
    audio = torch.zeros((1, samples), dtype=torch.float32)
    lengths = torch.tensor([samples], device=target, dtype=torch.long)
    started = time.perf_counter()
    codes = codec.encode(audio, lengths)[0][0].cpu()
    return {
        "input": "digital_silence",
        "seconds": seconds,
        "samples": samples,
        "code_frames": int(codes.shape[-1]),
        "elapsed_seconds": time.perf_counter() - started,
    }


class StagedReferenceCodec(DAC):
    """DAC with decode modules resident and reference-only modules staged."""

    @property
    def device(self) -> torch.device:
        return self._decode_device

    def configure_reference_staging(
        self,
        decode_device: str | torch.device,
        offload_device: str | torch.device = "cpu",
    ) -> None:
        self._decode_device = torch.device(decode_device)
        self._reference_offload_device = torch.device(offload_device)
        self._reference_lock = threading.Lock()

        # These modules are used by ``from_indices`` and remain resident.
        self.quantizer.semantic_quantizer.to(self._decode_device)
        self.quantizer.quantizer.to(self._decode_device)
        self.quantizer.post_module.to(self._decode_device)
        self.quantizer.upsample.to(self._decode_device)
        self.decoder.to(self._decode_device)

        # These modules are required only while a new reference is encoded.
        self.encoder.to(self._reference_offload_device)
        self.quantizer.downsample.to(self._reference_offload_device)
        self.quantizer.pre_module.to(self._reference_offload_device)

    @torch.inference_mode()
    def encode(
        self,
        audio_data: torch.Tensor,
        audio_lengths: torch.Tensor | None = None,
        n_quantizers: int | None = None,
        **kwargs,
    ):
        """Encode reference codes, staging only the required modules on CUDA."""

        if not hasattr(self, "_reference_lock"):
            return super().encode(
                audio_data,
                audio_lengths=audio_lengths,
                n_quantizers=n_quantizers,
                **kwargs,
            )

        with self._reference_lock:
            reference_modules = (
                self.encoder,
                self.quantizer.downsample,
                self.quantizer.pre_module,
            )
            for module in reference_modules:
                module.to(self._decode_device)

            try:
                dtype = next(self.encoder.parameters()).dtype
                audio_data = audio_data.to(device=self._decode_device, dtype=dtype)
                if audio_data.ndim == 2:
                    audio_data = audio_data.unsqueeze(1)
                length = audio_data.shape[-1]
                right_pad = (
                    math.ceil(length / self.frame_length) * self.frame_length - length
                )
                audio_data = torch.nn.functional.pad(audio_data, (0, right_pad))
                if audio_lengths is None:
                    audio_lengths = torch.tensor(
                        [length + right_pad],
                        device=self._decode_device,
                        dtype=torch.long,
                    )
                else:
                    audio_lengths = audio_lengths.to(self._decode_device)

                z = self.encoder(audio_data)
                z = self.quantizer.downsample(z)
                z = self.quantizer.pre_module(z)
                semantic_z, semantic_codes, *_ = self.quantizer.semantic_quantizer(z)
                residual_z = z - semantic_z
                _, residual_codes, *_ = self.quantizer.quantizer(
                    residual_z,
                    n_quantizers=n_quantizers,
                )
                indices = torch.cat([semantic_codes, residual_codes], dim=1)
                indices_lens = torch.ceil(audio_lengths / self.frame_length).long()
            finally:
                if self._decode_device.type == "cuda":
                    torch.cuda.synchronize(self._decode_device)
                for module in reference_modules:
                    module.to(self._reference_offload_device)
                if self._decode_device.type == "cuda":
                    torch.cuda.empty_cache()

            return indices, indices_lens


def _tensor_bytes(tensor: torch.Tensor | None) -> int:
    if tensor is None:
        return 0
    return tensor.numel() * tensor.element_size()


@torch.inference_mode()
def compact_codec_inference_buffers(codec: torch.nn.Module) -> dict[str, Any]:
    """Remove dead causal masks and bound RoPE tables to configured limits.

    ``WindowLimitedTransformer.forward`` always constructs an exact mask for
    the current input and passes it to its parent implementation. Therefore the
    inherited 32768-square causal mask is not read on this path. Its RoPE table
    is used, but the configured block size is the model's supported inference
    limit and is far smaller than the inherited 327680-frame table.
    """

    from fish_speech.models.dac.modded_dac import WindowLimitedTransformer

    records: list[dict[str, Any]] = []
    saved_bytes = 0
    for name, module in codec.named_modules():
        if not isinstance(module, WindowLimitedTransformer):
            continue

        causal_mask = module.causal_mask
        freqs_cis = module.freqs_cis
        if freqs_cis is None:
            raise RuntimeError(f"Codec transformer {name} has no RoPE table")

        frame_limit = int(module.config.block_size)
        if frame_limit <= 0 or frame_limit > freqs_cis.shape[0]:
            raise RuntimeError(
                f"Invalid codec RoPE limit for {name}: {frame_limit} "
                f"of {freqs_cis.shape[0]}"
            )

        before_mask_bytes = _tensor_bytes(causal_mask)
        before_rope_bytes = _tensor_bytes(freqs_cis)
        device = freqs_cis.device
        module.causal_mask = torch.empty(0, dtype=torch.bool, device=device)
        module.freqs_cis = freqs_cis[:frame_limit].clone()
        after_rope_bytes = _tensor_bytes(module.freqs_cis)
        module._compact_inference_frame_limit = frame_limit

        records.append(
            {
                "module": name,
                "frame_limit": frame_limit,
                "removed_causal_mask_bytes": before_mask_bytes,
                "rope_bytes_before": before_rope_bytes,
                "rope_bytes_after": after_rope_bytes,
            }
        )
        saved_bytes += before_mask_bytes + before_rope_bytes - after_rope_bytes

    if len(records) != 3:
        raise RuntimeError(
            f"Expected three window-limited codec transformers, found {len(records)}"
        )

    report = {
        "policy": "compact_windowed_inference_buffers",
        "windowed_transformers": len(records),
        "theoretical_saved_bytes": saved_bytes,
        "records": records,
    }
    codec._compact_inference_buffers_report = report
    return report


@torch.inference_mode()
def load_compact_codec_model(
    config_name: str,
    checkpoint_path: str | Path,
    device: str | torch.device = "cuda:0",
    precision: torch.dtype = torch.bfloat16,
    offload_reference: bool = False,
) -> torch.nn.Module:
    """Load the pinned codec with compact buffers before CUDA placement."""

    from hydra.utils import instantiate
    from omegaconf import OmegaConf

    from fish_speech.models.dac import modded_dac as modded_dac_module

    config_path = (
        Path(modded_dac_module.__file__).resolve().parents[2]
        / "configs"
        / f"{config_name}.yaml"
    )
    cfg = OmegaConf.load(config_path)
    if offload_reference:
        cfg._target_ = "experimental.codec.StagedReferenceCodec"

    codec = instantiate(cfg)
    state_dict = torch.load(
        checkpoint_path,
        map_location="cpu",
        mmap=True,
        weights_only=True,
    )
    if "state_dict" in state_dict:
        state_dict = state_dict["state_dict"]
    if any("generator" in key for key in state_dict):
        state_dict = {
            key.replace("generator.", ""): value
            for key, value in state_dict.items()
            if "generator." in key
        }

    load_result = codec.load_state_dict(state_dict, strict=False, assign=True)
    unexpected = [
        key
        for key in load_result.unexpected_keys
        if not key.endswith(("causal_mask", "freqs_cis"))
    ]
    if load_result.missing_keys or unexpected:
        raise RuntimeError(
            "Unexpected compact codec checkpoint mismatch: "
            f"missing={load_result.missing_keys[:5]}, unexpected={unexpected[:5]}"
        )

    report = compact_codec_inference_buffers(codec)
    codec.eval()
    codec.to(dtype=precision)
    if offload_reference:
        codec.configure_reference_staging(device)
        report["reference_path"] = "staged_from_cpu_to_decode_device"
    else:
        codec.to(device=device)
        report["reference_path"] = "resident_on_decode_device"
    codec._compact_inference_buffers_report = report
    del state_dict
    gc.collect()
    if torch.cuda.is_available() and torch.device(device).type == "cuda":
        torch.cuda.empty_cache()
    return codec


__all__ = [
    "StagedReferenceCodec",
    "compact_codec_inference_buffers",
    "load_compact_codec_model",
    "load_reference_audio_soundfile",
    "warm_reference_encoder",
]