import os
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # MUST be first (before torch)
import torch
import torch.nn.functional as F
import torchaudio
import numpy as np
import math
import tempfile
import soundfile as sf
from pathlib import Path
from typing import Optional, Tuple

import gradio as gr
from transformers import AutoModel, Wav2Vec2FeatureExtractor, Wav2Vec2Model
from huggingface_hub import hf_hub_download, snapshot_download

# ─── Model source code (vendored from QuarkAudio-UniSE repo) ───
from model.bicodec import BiCodecTokenizer
from model.llm import LLM_SFT

# ─── Config ───
SPARK_TTS_REPO = "SparkAudio/Spark-TTS-0.5B"
UNISE_REPO = "QuarkAudio/QuarkAudio-UniSE"
UNISE_CKPT = "epoch=20-step=109367.ckpt"

STFT_CONFIG = {
    "hop_length": 320,
    "win_length": 640,
    "n_fft": 640,
    "n_mels": 80,
}

LLM_CONFIG = {
    "num_tasks": 3,
    "task_map": {"se": 0, "tse": 1, "rtse": 2},
    "feats_dim": 768,
    "llm_base_config": {
        "cond_dim": 80,
        "global_size": 4096,
        "semantic_size": 8192,
        "hidden_size": 512,
        "num_layers": 12,
        "num_attention_heads": 8,
        "dropout_p": 0.1,
        "max_position_embeddings": 4096,
        "label_smoothing": 0.1,
        "conformer_params": {
            "num_layers": 6,
            "dim": 512,
            "heads": 8,
            "dim_head": 64,
            "depthwise_conv_kernel_size": 31,
            "ff_mult": 4,
            "dropout": 0.1,
            "qk_norm": None,
            "pe_attn_head": 1,
        },
    },
}

SAMPLE_RATE = 16000
SEG_LEN = 5 * SAMPLE_RATE  # 5 seconds per segment

# ─── Load models at module scope ───

# 1. Download BiCodec checkpoints
print("Downloading BiCodec and wav2vec2 checkpoints...")
bicodec_dir = snapshot_download(SPARK_TTS_REPO, repo_type="model")
bicodec_ckpt_dir = Path(bicodec_dir)
print(f"BiCodec dir: {bicodec_ckpt_dir}")

# 2. Load BiCodec tokenizer (loads BiCodec model + wav2vec2 feature extractor)
print("Loading BiCodec tokenizer...")
tokenizer = BiCodecTokenizer(model_dir=str(bicodec_ckpt_dir))
tokenizer = tokenizer.to("cuda")

# 3. Load WavLM for semantic feature extraction
print("Loading WavLM...")
semantic_model = AutoModel.from_pretrained("microsoft/wavlm-base-plus").to("cuda").eval()
semantic_model.requires_grad_(False)

# 4. Build LLM and load checkpoint
print("Loading UniSE LLM...")
llm = LLM_SFT(**LLM_CONFIG)

unise_ckpt_path = hf_hub_download(UNISE_REPO, UNISE_CKPT, repo_type="model")
ckpt = torch.load(unise_ckpt_path, map_location="cpu", weights_only=False)
state_dict = ckpt["state_dict"]

# Keys are prefixed with "dnn." — strip it
clean_sd = {}
for k, v in state_dict.items():
    if k.startswith("dnn."):
        clean_sd[k[4:]] = v
    else:
        clean_sd[k] = v

missing, unexpected = llm.load_state_dict(clean_sd, strict=False)
if missing:
    print(f"Missing keys: {missing}")
if unexpected:
    print(f"Unexpected keys: {unexpected}")

llm = llm.to("cuda").eval()

print("All models loaded successfully!")


# ─── Feature extraction helpers (from model.py) ───

@torch.no_grad()
def extract_semantic_features(wavs: torch.Tensor) -> torch.Tensor:
    """Extract WavLM features by averaging across all hidden layers."""
    wavs = F.pad(wavs, (160, 160))
    feats = semantic_model(wavs, output_hidden_states=True)
    feats_mix = torch.stack(feats.hidden_states, dim=1).mean(1)
    return feats_mix.detach()


@torch.no_grad()
def stft_logmel(x: torch.Tensor) -> torch.Tensor:
    """Compute log-mel spectrogram."""
    assert x.ndim == 2
    hop_length = STFT_CONFIG["hop_length"]
    win_length = STFT_CONFIG["win_length"]
    n_fft = STFT_CONFIG["n_fft"]
    n_mels = STFT_CONFIG["n_mels"]

    pad_length = math.ceil(x.size(-1) / hop_length) * hop_length - x.size(-1)
    x = F.pad(
        x,
        (
            (win_length - hop_length) // 2,
            pad_length + (win_length - hop_length) // 2,
        ),
    )
    spec = torch.stft(
        x,
        n_fft,
        hop_length,
        win_length=win_length,
        window=torch.hann_window(win_length).to(x.device),
        onesided=True,
        center=False,
        return_complex=True,
    ).transpose(1, 2)

    if not hasattr(stft_logmel, "_fb"):
        from torchaudio.functional import melscale_fbanks
        fb = melscale_fbanks(
            n_freqs=n_fft // 2 + 1,
            f_min=0.0,
            f_max=8000.0,
            n_mels=n_mels,
            sample_rate=SAMPLE_RATE,
        )
        stft_logmel._fb = fb.to(x.device)

    mag = spec.abs()
    mel = mag @ stft_logmel._fb
    mel = torch.log(mel + 1e-10)
    return mel


def _decode_audio(audio_path: str) -> Tuple[np.ndarray, int]:
    """Decode an audio file to a float32 numpy array + sample rate.

    Tries soundfile first (wav/flac/ogg); falls back to torchaudio's
    ffmpeg backend for compressed formats such as mp3/m4a/aac that
    libsndfile cannot decode.
    """
    try:
        data, sr = sf.read(audio_path, dtype="float32", always_2d=False)
        return np.asarray(data, dtype=np.float32), int(sr)
    except Exception:
        pass
    # Fallback 1: torchaudio (ffmpeg/soundfile backends) handles mp3/m4a/ogg/etc.
    try:
        wav, sr = torchaudio.load(audio_path)  # (channels, T), float32
        return wav.numpy().T, int(sr)  # -> (T, channels) to match sf.read layout
    except Exception:
        pass
    # Fallback 2: librosa (audioread/ffmpeg) as a last resort for compressed formats.
    import librosa
    data, sr = librosa.load(audio_path, sr=None, mono=False)
    if data.ndim > 1:
        data = data.T  # librosa returns (channels, T); match sf.read (T, channels)
    return np.asarray(data, dtype=np.float32), int(sr)


@torch.no_grad()
def load_and_resample(audio_path: str, target_sr: int = SAMPLE_RATE) -> torch.Tensor:
    """Load an audio file of any common format (wav/mp3/flac/ogg/m4a) or mic
    recording, downmix to mono, resample to 16kHz, and peak-normalize.

    Normalization + resampling make the native Gradio mic input (typically
    44.1/48kHz and often quiet) usable without error.
    """
    import soxr

    data, sr = _decode_audio(audio_path)

    # Downmix to mono
    if data.ndim > 1:
        data = data.mean(axis=1)
    data = data.astype(np.float32)

    # Resample to the model sample rate
    if sr != target_sr:
        data = soxr.resample(data, sr, target_sr, quality="VHQ")

    # Peak-normalize (guards against very quiet mic input / clipping)
    peak = float(np.max(np.abs(data))) if data.size else 0.0
    if peak > 1e-6:
        data = data / peak * 0.99

    return torch.from_numpy(np.ascontiguousarray(data)).float()


@spaces.GPU
def enhance_speech(
    input_audio: str,
    task_mode: str = "Speech Enhancement (SE)",
    enroll_audio: Optional[str] = None,
) -> str:
    """Enhance noisy speech audio using the UniSE model.

    Args:
        input_audio: Path to the noisy/mixed audio file.
        task_mode: One of 'Speech Enhancement (SE)', 'Target Speaker Extraction (TSE)', 'Speech Separation (SS)'.
        enroll_audio: Path to reference enrollment audio (required for TSE, optional for SE).
    """
    if input_audio is None:
        return None

    # Determine task from mode
    if "TSE" in task_mode:
        mode = "tse"
    elif "SS" in task_mode:
        mode = "ss"
    else:
        mode = "se"

    # Load input audio
    src = load_and_resample(input_audio).to("cuda").unsqueeze(0)  # (1, T)

    # Load enrollment if provided (for TSE)
    enroll = None
    if enroll_audio is not None and mode in ("tse", "rtse"):
        enroll = load_and_resample(enroll_audio).to("cuda").unsqueeze(0)  # (1, T)

    do_sample = False  # greedy for deterministic output

    if mode == "se":
        # Speech Enhancement: segment into 5s chunks, process each
        pad_len = math.ceil(src.size(-1) / SEG_LEN) * SEG_LEN - src.size(-1)
        seg_src = np.pad(src.cpu().numpy(), [(0, 0), (0, pad_len)], "wrap")
        seg_src = torch.from_numpy(seg_src).to(src.device)
        seg_src = seg_src.reshape(-1, SEG_LEN)
        seg_src = seg_src / src.abs().max(dim=-1, keepdim=True)[0]

        mix_mel = stft_logmel(seg_src)
        mix_feats = extract_semantic_features(seg_src)

        global_ids, semantic_ids = llm.generate(
            task_name="se",
            enroll_mel=None,
            enroll_feats=None,
            mix_mel=mix_mel,
            mix_feats=mix_feats,
            do_sample=do_sample,
        )

        est = tokenizer.detokenize(global_ids.unsqueeze(1), semantic_ids).squeeze(1)
        est = est.reshape(-1)[:src.size(-1)]
        est = est.cpu().numpy()

    elif mode == "tse":
        if enroll is None:
            return "Error: Target Speaker Extraction requires enrollment audio."

        seg_len = SEG_LEN
        pad_len = math.ceil(src.size(-1) / seg_len) * seg_len - src.size(-1)
        seg_src = np.pad(src.cpu().numpy(), [(0, 0), (0, pad_len)], "wrap")
        seg_src = torch.from_numpy(seg_src).to(src.device)
        seg_src = seg_src.reshape(-1, seg_len)

        enroll_mel = stft_logmel(enroll)
        enroll_feats = extract_semantic_features(enroll)
        enroll_mel = torch.cat([enroll_mel for _ in range(seg_src.size(0))], dim=0)
        enroll_feats = torch.cat([enroll_feats for _ in range(seg_src.size(0))], dim=0)

        mix_mel = stft_logmel(seg_src)
        mix_feats = extract_semantic_features(seg_src)

        global_ids, semantic_ids = llm.generate(
            task_name="tse",
            enroll_mel=enroll_mel,
            enroll_feats=enroll_feats,
            mix_mel=mix_mel,
            mix_feats=mix_feats,
            do_sample=do_sample,
        )

        est = tokenizer.detokenize(global_ids.unsqueeze(1), semantic_ids).squeeze(1)
        est = est.reshape(-1)[:src.size(-1)]
        est = est.cpu().numpy()

    elif mode == "ss":
        # Speech Separation: first SE on first 5s to get enrollment,
        # then TSE to extract speaker 1, then rTSE for speaker 2
        seg_len = SEG_LEN
        if src.size(-1) > seg_len:
            seg_src = src[:, :seg_len]
        else:
            seg_src = np.pad(src.cpu().numpy(), [(0, 0), (0, seg_len - src.size(-1))], "wrap")
            seg_src = torch.from_numpy(seg_src).to(src.device)

        mix_mel = stft_logmel(seg_src)
        mix_feats = extract_semantic_features(seg_src)

        global_ids, semantic_ids = llm.generate(
            task_name="se",
            enroll_mel=None,
            enroll_feats=None,
            mix_mel=mix_mel,
            mix_feats=mix_feats,
            do_sample=do_sample,
        )

        enroll = tokenizer.detokenize(global_ids.unsqueeze(1), semantic_ids).squeeze(1)
        enroll = enroll[:, :seg_len]
        enroll = enroll / (torch.max(torch.abs(enroll)) + 1e-5) * 0.99

        enroll_mel = stft_logmel(enroll)
        enroll_feats = extract_semantic_features(enroll)

        pad_len = math.ceil(src.size(-1) / seg_len) * seg_len - src.size(-1)
        seg_src_full = np.pad(src.cpu().numpy(), [(0, 0), (0, pad_len)], "wrap")
        seg_src_full = torch.from_numpy(seg_src_full).to(src.device)
        seg_src_full = seg_src_full.reshape(-1, seg_len)

        enroll_mel = torch.cat([enroll_mel for _ in range(seg_src_full.size(0))], dim=0)
        enroll_feats = torch.cat([enroll_feats for _ in range(seg_src_full.size(0))], dim=0)
        mix_mel_full = stft_logmel(seg_src_full)
        mix_feats_full = extract_semantic_features(seg_src_full)

        # Extract speaker 1
        global_ids, semantic_ids = llm.generate(
            task_name="tse",
            enroll_mel=enroll_mel,
            enroll_feats=enroll_feats,
            mix_mel=mix_mel_full,
            mix_feats=mix_feats_full,
            do_sample=do_sample,
        )
        est_s1 = tokenizer.detokenize(global_ids.unsqueeze(1), semantic_ids).squeeze(1)
        est_s1 = est_s1.reshape(-1)[:src.size(-1)].cpu().numpy()

        # Extract speaker 2 (rTSE)
        global_ids, semantic_ids = llm.generate(
            task_name="rtse",
            enroll_mel=enroll_mel,
            enroll_feats=enroll_feats,
            mix_mel=mix_mel_full,
            mix_feats=mix_feats_full,
            do_sample=do_sample,
        )
        est_s2 = tokenizer.detokenize(global_ids.unsqueeze(1), semantic_ids).squeeze(1)
        est_s2 = est_s2.reshape(-1)[:src.size(-1)].cpu().numpy()

        # Save both speakers as a stereo file
        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
            sf.write(f.name, np.stack([est_s1, est_s2], axis=-1), SAMPLE_RATE)
            return f.name

    # Save output for SE and TSE
    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
        sf.write(f.name, est, SAMPLE_RATE)
        return f.name


# ─── Gradio UI ───

CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("# UniSE: Unified Speech Enhancement")
        gr.Markdown(
            "Upload noisy or mixed speech audio and get enhanced output. "
            "Supports **Speech Enhancement** (denoising), **Target Speaker Extraction** "
            "(extract a target speaker using reference audio), and **Speech Separation** "
            "(separate mixed speakers into individual streams)."
        )
        gr.Markdown(
            "📄 [Paper](https://arxiv.org/abs/2510.20441) | "
            "🤗 [Model](https://huggingface.co/QuarkAudio/QuarkAudio-UniSE) | "
            "💻 [Code](https://github.com/alibaba/unified-audio/tree/main/QuarkAudio-UniSE)"
        )

        with gr.Row():
            with gr.Column(scale=1):
                task_selector = gr.Radio(
                    choices=[
                        "Speech Enhancement (SE)",
                        "Target Speaker Extraction (TSE)",
                        "Speech Separation (SS)",
                    ],
                    value="Speech Enhancement (SE)",
                    label="Task Mode",
                )
                input_audio = gr.Audio(
                    label="Input Audio (noisy/mixed)",
                    type="filepath",
                )
                enroll_audio = gr.Audio(
                    label="Reference Enrollment Audio (required for TSE, ignored for SE/SS)",
                    type="filepath",
                    visible=True,
                )
                run_btn = gr.Button("Enhance", variant="primary")

            with gr.Column(scale=1):
                output_audio = gr.Audio(label="Enhanced Output", type="filepath")

        gr.Examples(
            examples=[
                ["examples/sr_noisy_1.wav", "Speech Enhancement (SE)", None],
                ["examples/sr_noisy_2.wav", "Speech Enhancement (SE)", None],
                ["examples/tse_noisy_1.wav", "Target Speaker Extraction (TSE)", "examples/tse_ref_1.wav"],
                ["examples/tse_noisy_2.wav", "Target Speaker Extraction (TSE)", "examples/tse_ref_2.wav"],
                ["examples/ss_noisy_1.wav", "Speech Separation (SS)", None],
            ],
            inputs=[input_audio, task_selector, enroll_audio],
            outputs=output_audio,
            fn=enhance_speech,
            cache_examples=True,
            cache_mode="lazy",
        )

    run_btn.click(
        fn=enhance_speech,
        inputs=[input_audio, task_selector, enroll_audio],
        outputs=output_audio,
        api_name="enhance",
    )

demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)