import sys

sys.stdout.reconfigure(line_buffering=True)

try:
    import spaces
except ImportError:
    # keep @spaces.GPU usable as a no-op; ZeroGPU requires this exact name.
    class spaces:
        class GPU:
            def __init__(self, func=None, duration=60):
                self.func = func

            def __call__(self, *args, **kwargs):
                if self.func is not None:
                    return self.func(*args, **kwargs)
                func = args[0]
                return func


import gradio as gr
import soundfile as sf
import torch
import torchaudio.functional as AF
from pyharp import ModelCard, build_endpoint

from model.speech_models.metricgan_generator import MetricGANGenerator
from model.speech_models.fullsubnet import PhaseInvariantFullSubNet

# Constructed directly rather than via the original repo's LightningCLI config loader,
# since only speech_model (never reverb_model/joint_loss_module) is used at inference.

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

PRESETS = {
    "Standard": {
        "cls": MetricGANGenerator,
        "ckpt": "checkpoints/unsupervised_bilstm/model.ckpt",
    },
    "Enhanced (phase-aware, weak supervision)": {
        "cls": PhaseInvariantFullSubNet,
        "ckpt": "checkpoints/phaseinv_fsn/model.ckpt",
    },
}

_models = {}  # preset name -> model instance already resident on DEVICE


def get_model(preset):
    """Builds and caches a preset's model."""
    if preset not in _models:
        cfg = PRESETS[preset]
        model = cfg["cls"](metrics=[])  # config's metrics are for validation logging, unused at inference
        model.load_state_dict_from_joint_model(cfg["ckpt"])
        model.eval()  # skips the original's .freeze(); @torch.inference_mode() already covers it
        # .to() recurses into FirstLevelModule's shared stft/istft singleton too
        model = model.to(DEVICE)
        _models[preset] = model
    return _models[preset]


model_card = ModelCard(
    name="U-DREAM",
    description="Removes room reverb from a speech recording. Mono only; stereo input is downmixed.",
    author="Louis Bahrman, Marius Rodrigues, Mathieu Fontaine, Gaël Richard",
    tags=["dereverberation", "speech enhancement"],
)


@spaces.GPU
@torch.inference_mode()
def process_fn(input_audio_path: str, preset: str) -> str:
    """Runs the selected dereverberation model on one audio file."""
    model = get_model(preset)

    data, fs = sf.read(input_audio_path, dtype="float32", always_2d=True)
    y = torch.from_numpy(data.T).to(DEVICE)  # (channels, samples)
    if y.size(0) > 1:
        y = y.mean(dim=0, keepdim=True)  # model is monaural-only (per paper)
    if fs != model.fs:
        # musicians shouldn't have to pre-convert sample rate themselves
        y = AF.resample(y, orig_freq=fs, new_freq=model.fs)
        fs = model.fs

    pred = model(y[None, ...])
    s = model.get_time(pred)
    s = s / s.abs().max()

    output_audio_path = input_audio_path.rsplit(".", 1)[0] + "_predicted_dry.wav"
    out = s[0].cpu().numpy()  # (1, samples): mono in guarantees mono out
    sf.write(output_audio_path, out[0], fs)
    return output_audio_path


with gr.Blocks() as demo:
    input_components = [
        gr.Audio(type="filepath", label="Input Audio").harp_required(True),
        gr.Dropdown(
            choices=[("Standard", "Standard"), ("Enhanced", "Enhanced (phase-aware, weak supervision)")],
            value="Standard",
            label="Preset",
            info="Standard: general-purpose reverb removal (paper's recommended BiLSTM default). "
            "Enhanced: can be better or worse (phase-invariant FullSubNet).",
        ),
    ]
    output_components = [
        gr.Audio(type="filepath", label="Output Audio").set_info("Dereverberated (dry) speech."),
    ]

    build_endpoint(
        model_card=model_card,
        input_components=input_components,
        output_components=output_components,
        process_fn=process_fn,
    )

if __name__ == "__main__":
    demo.queue().launch(pwa=True)