import functools
import tempfile
import wave
from typing import Any

try:
    import spaces
except ImportError:
    class _SpacesFallback:
        @staticmethod
        def GPU(*args: Any, **kwargs: Any) -> Any:
            def decorator(fn: Any) -> Any:
                return fn

            return decorator

    spaces = _SpacesFallback()

import gradio as gr
from gradio import processing_utils
import torch
import torchaudio
import torchaudio.functional as F
from transformers import AutoModel


MODEL_IDS = {
    "MOSS-Audio-Tokenizer-v2": "OpenMOSS-Team/MOSS-Audio-Tokenizer-v2",
    "MOSS-Audio-Tokenizer": "OpenMOSS-Team/MOSS-Audio-Tokenizer",
    "MOSS-Audio-Tokenizer-Nano": "OpenMOSS-Team/MOSS-Audio-Tokenizer-Nano",
}

DEFAULT_NQ = 8
MAX_NQ = 32


def _get_nested_attr(obj: Any, *names: str) -> Any:
    for name in names:
        if obj is None:
            continue
        if hasattr(obj, name):
            value = getattr(obj, name)
            if value is not None:
                return value
        if isinstance(obj, dict) and obj.get(name) is not None:
            return obj[name]
    return None


def _model_config(model: torch.nn.Module) -> Any:
    return getattr(model, "config", None)


def _target_sample_rate(model: torch.nn.Module) -> int:
    config = _model_config(model)
    sample_rate = _get_nested_attr(
        config,
        "sample_rate",
        "sampling_rate",
        "audio_sample_rate",
        "input_sample_rate",
    )
    if sample_rate is None:
        sample_rate = _get_nested_attr(
            _get_nested_attr(config, "audio_encoder", "codec_config", "model_config"),
            "sample_rate",
            "sampling_rate",
        )
    return int(sample_rate or 24000)


def _target_channels(model: torch.nn.Module) -> int:
    config = _model_config(model)
    channels = _get_nested_attr(
        config,
        "channels",
        "audio_channels",
        "input_channels",
        "num_channels",
        "number_channels",
    )
    if channels is None:
        channels = _get_nested_attr(
            _get_nested_attr(config, "audio_encoder", "codec_config", "model_config"),
            "channels",
            "audio_channels",
            "input_channels",
            "number_channels",
        )
    return int(channels or 1)


def _codebook_count(model: torch.nn.Module) -> int:
    config = _model_config(model)
    candidates = (
        config,
        _get_nested_attr(config, "quantizer", "rvq", "codec_config", "model_config"),
        getattr(model, "quantizer", None),
    )
    for candidate in candidates:
        count = _get_nested_attr(
            candidate,
            "nq",
            "n_q",
            "num_quantizers",
            "num_codebooks",
            "codebook_num",
        )
        if count is not None:
            return int(count)
    return MAX_NQ


@functools.lru_cache(maxsize=1)
def load_model(model_id: str) -> torch.nn.Module:
    model = AutoModel.from_pretrained(model_id, trust_remote_code=True)
    model.eval()
    return model


def _device() -> torch.device:
    return torch.device("cuda" if torch.cuda.is_available() else "cpu")


def preprocess_audio(audio_path: str, model: torch.nn.Module, device: torch.device) -> tuple[torch.Tensor, int]:
    source_sr, waveform = processing_utils.audio_from_file(audio_path)
    wav = torch.from_numpy(waveform).float()
    if wav.dim() == 1:
        wav = wav.unsqueeze(0)
    else:
        wav = wav.transpose(0, 1)

    if wav.abs().max() > 1.0:
        wav = wav / max(float(wav.abs().max()), 1.0)

    target_channels = _target_channels(model)
    if wav.shape[0] != target_channels:
        if target_channels == 1:
            wav = wav.mean(dim=0, keepdim=True)
        elif wav.shape[0] == 1:
            wav = wav.repeat(target_channels, 1)
        else:
            wav = wav[:target_channels]

    target_sr = _target_sample_rate(model)
    if source_sr != target_sr:
        wav = F.resample(wav, source_sr, target_sr)

    return wav.unsqueeze(0).to(device), target_sr


def _extract_audio_codes(encoded: Any) -> torch.Tensor:
    if hasattr(encoded, "audio_codes"):
        return encoded.audio_codes
    if isinstance(encoded, dict) and "audio_codes" in encoded:
        return encoded["audio_codes"]
    if isinstance(encoded, (tuple, list)) and encoded:
        return encoded[0]
    raise RuntimeError("Model encode output does not contain `audio_codes`.")


def _extract_decoded_audio(decoded: Any) -> torch.Tensor:
    if hasattr(decoded, "audio_values"):
        return decoded.audio_values
    if hasattr(decoded, "audio"):
        return decoded.audio
    if isinstance(decoded, dict):
        for key in ("audio_values", "audio", "wav", "waveform"):
            if key in decoded:
                return decoded[key]
    if torch.is_tensor(decoded):
        return decoded
    if isinstance(decoded, (tuple, list)) and decoded:
        return decoded[0]
    raise RuntimeError("Model decode output does not contain reconstructed audio.")


def _codes_shape(codes: torch.Tensor) -> str:
    return " x ".join(str(dim) for dim in codes.shape)


def _slice_rvq_prefix(codes: torch.Tensor, nq: int, batch_size: int) -> torch.Tensor:
    if codes.dim() == 3:
        if codes.shape[1] == batch_size:
            return codes[:nq]
        if codes.shape[0] == batch_size:
            return codes[:, :nq, :]
    if codes.dim() == 2:
        return codes[:nq]
    return codes[:nq]


def _to_audio_file(audio: torch.Tensor, sample_rate: int) -> str:
    audio = audio.detach().cpu().float()
    while audio.dim() > 2:
        audio = audio.squeeze(0)
    if audio.dim() == 1:
        audio = audio.unsqueeze(0)
    if audio.dim() == 2 and audio.shape[0] > audio.shape[1]:
        audio = audio.transpose(0, 1)
    audio = audio.clamp(-1.0, 1.0)

    output_file = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
    output_file.close()
    audio_np = (audio.transpose(0, 1).numpy() * 32767.0).astype("<i2")
    with wave.open(output_file.name, "wb") as wav_file:
        wav_file.setnchannels(audio_np.shape[1])
        wav_file.setsampwidth(2)
        wav_file.setframerate(sample_rate)
        wav_file.writeframes(audio_np.tobytes())
    return output_file.name


@spaces.GPU(duration=120)
def reconstruct(audio_path: str | None, model_name: str, nq: int) -> tuple[str | None, str]:
    if audio_path is None:
        raise gr.Error("请先上传或录制一段音频。")

    model_id = MODEL_IDS[model_name]
    device = _device()
    model = load_model(model_id)
    model.to(device)
    wav, sample_rate = preprocess_audio(audio_path, model, device)

    max_nq = _codebook_count(model)
    nq = int(max(1, min(nq, max_nq)))

    with torch.inference_mode():
        encoded = model.encode(wav, return_dict=True)
        audio_codes = _extract_audio_codes(encoded)
        sliced_codes = _slice_rvq_prefix(audio_codes, nq, wav.shape[0])
        decoded = model.decode(sliced_codes, return_dict=True)
        reconstructed = _extract_decoded_audio(decoded)

    info = (
        f"模型: {model_id}\n"
        f"输入采样率: {sample_rate} Hz\n"
        f"模型可用 RVQ 层数: {max_nq}\n"
        f"本次使用 nq: {nq}\n"
        f"RVQ token shape: {_codes_shape(audio_codes)}"
    )
    return _to_audio_file(reconstructed, sample_rate), info


with gr.Blocks(title="MOSS Audio Tokenizer Reconstruction") as demo:
    gr.Markdown("# MOSS-Audio-Tokenizer 系列音频重建")

    with gr.Row():
        with gr.Column(scale=1):
            model_input = gr.Dropdown(
                choices=list(MODEL_IDS.keys()),
                value="MOSS-Audio-Tokenizer-v2",
                label="模型",
            )
            audio_input = gr.Audio(
                sources=["upload", "microphone"],
                type="filepath",
                label="输入音频",
            )
            nq_input = gr.Slider(
                minimum=1,
                maximum=MAX_NQ,
                value=DEFAULT_NQ,
                step=1,
                label="期待的 nq",
            )
            run_button = gr.Button("重建音频", variant="primary")

        with gr.Column(scale=1):
            audio_output = gr.Audio(label="重建音频")
            info_output = gr.Textbox(label="编码信息", lines=6)

    run_button.click(
        fn=reconstruct,
        inputs=[audio_input, model_input, nq_input],
        outputs=[audio_output, info_output],
    )


if __name__ == "__main__":
    demo.launch()