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

import spaces
import torch
import numpy as np
import soundfile as sf
import tempfile
import time
from huggingface_hub import hf_hub_download

from models.gap_urgenet import GAP_URGENet

REPO_ID = "Xiaobin-Rong/gap-urgenet"


def _download_checkpoints():
    """Download all 5 checkpoints from HuggingFace Hub."""
    ckpts = {}
    for fname in ["DeWavLM-Omni.pt", "Adapter.pt", "Vocoder.pt",
                  "Predictor.pt", "PostNet.pt"]:
        ckpts[fname] = hf_hub_download(repo_id=REPO_ID, filename=fname)
    return ckpts


print("[*] Downloading checkpoints from HuggingFace Hub...")
_ckpt_paths = _download_checkpoints()
print("[*] All checkpoints downloaded.")

print("[*] Loading GAP-URGENet model...")
# Monkey-patch torch.load to avoid weights_only default issue in torch 2.6+
_orig_load = torch.load
torch.load = lambda *a, **k: _orig_load(*a, **{**k, "weights_only": k.get("weights_only", False)})

model = GAP_URGENet(
    dewavlm_ckpt_path=_ckpt_paths["DeWavLM-Omni.pt"],
    adapter_ckpt_path=_ckpt_paths["Adapter.pt"],
    vocoder_ckpt_path=_ckpt_paths["Vocoder.pt"],
    predictor_ckpt_path=_ckpt_paths["Predictor.pt"],
    postnet_ckpt_path=_ckpt_paths["PostNet.pt"],
).to("cuda").eval()

# Restore original torch.load
torch.load = _orig_load
print("[*] Model loaded and moved to CUDA.")


@spaces.GPU(duration=60)
def enhance_audio(audio_path: str, enable_plc: bool = True) -> str:
    """Enhance noisy speech audio using GAP-URGENet.

    Args:
        audio_path: Path to the noisy audio file (wav, flac, etc.).
        enable_plc: Whether to perform packet loss concealment (PLC).

    Returns:
        Path to the enhanced audio file (WAV format).
    """
    t0 = time.perf_counter()

    audio, fs = sf.read(audio_path, dtype='float32')

    # Handle mono / stereo
    if audio.ndim > 1:
        audio = audio[:, 0]  # take first channel

    input_tensor = torch.FloatTensor(audio).unsqueeze(0).to("cuda")

    with torch.inference_mode():
        output = model(input_tensor, sr_in=fs, sr_out=fs, enable_plc=enable_plc)

    enhanced = output.cpu().detach().numpy().squeeze()

    # Normalize to preserve original scale
    scale = np.max(np.abs(audio))
    if scale > 0:
        enhanced = enhanced / (np.max(np.abs(enhanced)) + 1e-8) * scale

    # Save to temp file
    out_path = tempfile.NamedTemporaryFile(suffix=".wav", delete=False).name
    sf.write(out_path, enhanced, fs)

    elapsed = time.perf_counter() - t0
    print(f"[*] Inference completed in {elapsed:.2f}s")
    return out_path


import gradio as gr

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

with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
    gr.Markdown(
        "# GAP-URGENet: Universal Speech Enhancement\n"
        "1st place in the ICASSP 2026 URGENT Challenge objective evaluation. "
        "Upload noisy speech audio and get enhanced output."
    )

    with gr.Column(elem_id="col-container"):
        with gr.Row():
            audio_input = gr.Audio(
                label="Noisy Audio",
                type="filepath",
            )
            audio_output = gr.Audio(
                label="Enhanced Audio",
                type="filepath",
            )

        with gr.Row():
            plc_checkbox = gr.Checkbox(label="Enable Packet Loss Concealment (PLC)", value=True)
            run_btn = gr.Button("Enhance Audio", variant="primary")

    run_btn.click(
        fn=enhance_audio,
        inputs=[audio_input, plc_checkbox],
        outputs=audio_output,
        api_name="enhance",
    )

    gr.Examples(
        examples=[
            ["examples/noisy_sample_1.wav", True],
            ["examples/noisy_sample_2.flac", True],
            ["examples/noisy_sample_3.flac", True],
        ],
        inputs=[audio_input, plc_checkbox],
        outputs=audio_output,
        fn=enhance_audio,
        cache_examples=True,
        cache_mode="lazy",
    )

demo.launch(mcp_server=True)