import sys
sys.stdout.reconfigure(line_buffering=True)  # real-time logs in HF Spaces

try:
    import spaces
    def gpu_decorator(func): return spaces.GPU(func)
except ImportError:
    def gpu_decorator(func): return func

import contextlib
import tempfile
import threading
import traceback

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

# ---- Paths and device ----
CKPT_PATH = "pretrained/VCTK_16k_4s_time-190000.pt"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# ---- Inference config ----
ARGS = OmegaConf.load("config/inference.yaml")
OP_HP = OmegaConf.load("config/operator.yaml")

SAMPLE_RATE = ARGS.exp.sample_rate
AUDIO_LEN = ARGS.exp.audio_len


# ---- Model loading ----
sampler = None
model_loading = True
model_error = None
model_ready = False  # has the network been moved onto the GPU yet?

def load_model():
    global sampler, model_loading, model_error
    try:
        print(f"Loading checkpoint from {CKPT_PATH}...")

        # imported here (not at top) so the Gradio server starts before these load

        from networks.ncsnpp import NCSNppTime
        from diff_params.edm import EDM
        from utils.training_utils import load_state_dict
        from testing.EulerHeunSamplerDPS import EulerHeunSamplerDPS

        network_cfg = OmegaConf.load("config/network.yaml")
        # stft stays as OmegaConf — NCSNppTime uses dot access on it (stft_kwargs.n_fft)
        stft_cfg = network_cfg.pop("stft")
        # built on CPU — ZeroGPU only intercepts CUDA calls inside an
        # @spaces.GPU-decorated call, not from this background thread
        network = NCSNppTime(stft=stft_cfg, **OmegaConf.to_container(network_cfg))

        # load_state_dict tries multiple key strategies ('ema', 'model', etc.)
        # to handle checkpoints saved in different formats
        state_dict = torch.load(CKPT_PATH, map_location="cpu", weights_only=False)
        load_state_dict(state_dict, ema=network)
        network.eval()

        diff_params_cfg = OmegaConf.load("config/diff_params.yaml")
        # sde_hp stays as OmegaConf — EDM uses dot access on it (sde_hp.sigma_data)
        sde_hp = diff_params_cfg.pop("sde_hp")
        diff_params = EDM(sde_hp=sde_hp, **OmegaConf.to_container(diff_params_cfg))

        sampler = EulerHeunSamplerDPS(network, diff_params, ARGS)
        print("Model ready.")

    except Exception:
        model_error = traceback.format_exc()
        print(f"Error loading model:\n{model_error}")
    finally:
        model_loading = False

# Load in background so the Gradio server starts immediately
threading.Thread(target=load_model, daemon=True).start()


# ---- pyharp model card ----
model_card = ModelCard(
    name="BUDDy - Blind Dereverberation",
    description="Removes room reverberation from a speech recording. No room measurements needed — the model estimates the room acoustics automatically.",
    author="Lemercier, Moliner, Welker, Välimäki, Gerkmann (2024)",
    tags=["speech", "dereverberation", "effect removal"],
)


# ---- Inference ----
@gpu_decorator
def process_fn(input_audio_path: str, num_steps: int):
    global model_ready
    if model_loading:
        raise gr.Error("Model is still loading, please wait a moment and try again.")
    if sampler is None:
        raise gr.Error(f"Model failed to load: {model_error}")
    if not model_ready:
        sampler.model.to(DEVICE)  # only safe here, inside @spaces.GPU
        model_ready = True

    from testing.operators.subband_filtering import BlindSubbandFiltering

    # Update step count from slider — also update args so get_gamma() uses the right T
    sampler.T = num_steps

    # using soundfile directly — torchaudio.load/save need torchcodec, which isn't installed
    data, sr = sf.read(input_audio_path)
    waveform = torch.tensor(data.T if data.ndim > 1 else data[None]).float()  # (channels, samples)

    # resampling to 16kHz
    if sr != SAMPLE_RATE:
        waveform = torchaudio.functional.resample(waveform, sr, SAMPLE_RATE)
    
    # converting to mono
    if waveform.shape[0] > 1:
        waveform = waveform.mean(dim=0, keepdim=True)
    
    # trimming/padding to exactly AUDIO_LEN
    if waveform.shape[-1] > AUDIO_LEN:
        waveform = waveform[..., :AUDIO_LEN]
    elif waveform.shape[-1] < AUDIO_LEN:
        waveform = torch.nn.functional.pad(waveform, (0, AUDIO_LEN - waveform.shape[-1]))
    y = waveform.squeeze(0).to(DEVICE)  # (AUDIO_LEN,)

    # Normalize to match sigma_data of the training set (0.05)
    y = ARGS.tester.posterior_sampling.warm_initialization.scaling_factor * y / (y.std() + 1e-8)
    y = y.unsqueeze(0)  # (1, AUDIO_LEN) — sampler expects a batch dimension

    # tester.py line 147
    operator = BlindSubbandFiltering(OP_HP, sample_rate=SAMPLE_RATE)
    with torch.no_grad():
        operator.update_H(use_noise=True)

    # Run the DPS sampler. No torch.inference_mode() here because the DPS
    # likelihood gradient uses torch.autograd.grad() on intermediate tensors.
    with contextlib.nullcontext():
        pred = sampler.predict_conditional(y, operator, shape=(1, AUDIO_LEN), blind=True)

    pred = pred.detach().cpu()
    if pred.dim() > 1:
        pred = pred.squeeze(0)

    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
        out_path = f.name
    sf.write(out_path, pred.numpy(), SAMPLE_RATE)
    return out_path


# ---- Gradio UI ----
with gr.Blocks() as demo:
    input_components = [
        gr.Audio(type="filepath", label="Reverberant Audio").harp_required(True),
        gr.Slider(
            minimum=30,
            maximum=400,
            step=1,
            value=201,
            label="Processing Steps",
            info="More steps = higher quality but slower. Paper default: 201. Max: 400.",
        ),
    ]
    output_components = [
        gr.Audio(type="filepath", label="Dereverberated Audio").set_info(
            "Clean speech with room reverb removed."
        ),
    ]
    build_endpoint(
        model_card=model_card,
        input_components=input_components,
        output_components=output_components,
        process_fn=process_fn,
    )

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