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 tempfile

import gradio as gr
import soundfile as sf
import torch
import torchaudio
from huggingface_hub import hf_hub_download
from omegaconf import OmegaConf
from pyharp import ModelCard, build_endpoint

from model.config import CHECKPOINTS, CHECKPOINTS_REPO
from model.cqtdiff import Unet_CQT_oct_with_attention
from model.edm import EDM
from model.restoration import restore_audio

# one entry per architecture ("piano"/"singing"), holding everything that's
# lazily built/loaded once and reused across requests for that architecture
_state = {}


def get_arch_state(base_cfg, device):
    """Lazily builds (once per architecture) and caches the network and EDM diffusion parameters on `device`."""
    key = base_cfg.architecture
    if key not in _state:
        net = Unet_CQT_oct_with_attention(base_cfg, device)
        net.to(device)
        net.eval()
        _state[key] = {"network": net, "diff_params": EDM(base_cfg), "loaded_checkpoint": None, "ltas_ref": None}
    return _state[key]


def load_checkpoint(label, device):
    """Downloads (if not already cached locally) and loads the checkpoint for the selected voice/instrument, reusing the shared network for its architecture."""
    filename, base_cfg = CHECKPOINTS[label]
    state = get_arch_state(base_cfg, device)

    if state["loaded_checkpoint"] != filename:
        ckpt_path = hf_hub_download(repo_id=CHECKPOINTS_REPO, filename=filename)
        checkpoint = torch.load(ckpt_path, map_location="cpu", weights_only=False)
        state["network"].load_state_dict(checkpoint["ema"])
        state["ltas_ref"] = checkpoint["LTAS"]
        state["loaded_checkpoint"] = filename

    return state["network"], state["diff_params"], state["ltas_ref"], base_cfg


model_card = ModelCard(
    name="BABE-2",
    description="Restores degraded historical piano or singing-voice recordings with a diffusion-based generative equalizer.",
    author="Eloi Moliner, Maija Turunen, Filip Elvander, Vesa Välimäki",
    tags=["restoration", "equalizer", "diffusion", "historical recordings"],
)


@spaces.GPU(duration=180)
def process_fn(input_audio_path: str, checkpoint_label: str, steps: int, strength: float) -> str:
    """Restores the input recording using the selected voice/instrument model, diffusion step count, and restoration strength."""
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    net, diff_params, ltas_ref, base_cfg = load_checkpoint(checkpoint_label, device)
    ltas_ref = ltas_ref.to(device)

    cfg = OmegaConf.create(OmegaConf.to_container(base_cfg, resolve=True))
    cfg.tester.T = int(steps)
    cfg.tester.posterior_sampling.xi = float(strength)

    audio, sr = sf.read(input_audio_path)
    sig = torch.tensor(audio, dtype=torch.float32)
    if sig.dim() > 1:
        sig = sig.mean(dim=-1)
    if sr != cfg.exp.sample_rate:
        sig = torchaudio.functional.resample(sig, sr, cfg.exp.sample_rate)

    def on_progress(fraction, desc):
        print(f"[{fraction:.0%}] {desc}")

    restored = restore_audio(sig, cfg, net, diff_params, ltas_ref, device, on_progress=on_progress)

    out_path = tempfile.NamedTemporaryFile(suffix=".wav", delete=False).name
    sf.write(out_path, restored.cpu().numpy(), cfg.exp.sample_rate)
    return out_path


def update_strength_default(checkpoint_label):
    """Updates the Restoration Strength slider to the selected model's paper-recommended default."""
    _, base_cfg = CHECKPOINTS[checkpoint_label]
    return gr.update(value=float(base_cfg.tester.posterior_sampling.xi))


with gr.Blocks() as demo:
    checkpoint_dropdown = gr.Dropdown(
        choices=list(CHECKPOINTS.keys()),
        value="Piano (MAESTRO)",
        label="Voice / Instrument",
        info="Singer-specific models give the most faithful restoration when your recording actually resembles that singer.",
    )

    input_components = [
        gr.Audio(type="filepath", label="Input Audio").harp_required(True),
        checkpoint_dropdown,
        gr.Slider(minimum=10, maximum=100, step=1, value=51, label="Processing Steps",
                   info="More steps can improve quality at the cost of processing time."),
        gr.Slider(minimum=0.0, maximum=2.0, step=0.05, value=1.0, label="Restoration Strength",
                   info="How strongly the output is guided to match the input recording."),
    ]
    output_components = [
        gr.Audio(type="filepath", label="Restored Audio").set_info("The restored recording."),
    ]

    checkpoint_dropdown.change(fn=update_strength_default, inputs=checkpoint_dropdown, outputs=input_components[3])

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

demo.queue().launch(pwa=True)