import sys
sys.stdout.reconfigure(line_buffering=True)  # flush print() calls immediately so HF Space logs are live

import os
import tempfile
import threading
import traceback

import gradio as gr
import librosa
import torch
import soundfile as sf
from huggingface_hub import hf_hub_download
from pyharp import ModelCard, build_endpoint

from models.resunet import ResUNet30
from models.clap_encoder import CLAP_Encoder

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

# Model loading state — populated by the background thread.
ss_model = None
query_encoder = None
model_loading = True
model_error = None


def load_model():
    global ss_model, query_encoder, model_loading, model_error
    try:
        print("Downloading checkpoints...")
        main_ckpt_path = hf_hub_download(
            repo_id="Audio-AGI/AudioSep",
            repo_type="space",
            filename="checkpoint/audiosep_base_4M_steps.ckpt",
        )
        clap_ckpt_path = hf_hub_download(
            repo_id="Audio-AGI/AudioSep",
            repo_type="space",
            filename="checkpoint/music_speech_audioset_epoch_15_esc_89.98.pt",
        )

        query_encoder = CLAP_Encoder(pretrained_path=clap_ckpt_path).eval()
        _ss_model = ResUNet30(input_channels=1, output_channels=1, condition_size=512)

        # Load weights from the Lightning checkpoint — keys are prefixed with "ss_model."
        state_dict = torch.load(main_ckpt_path, map_location="cpu", weights_only=False)["state_dict"]
        weights = {
            k.removeprefix("ss_model."): v
            for k, v in state_dict.items()
            if k.startswith("ss_model.")
        }
        _ss_model.load_state_dict(weights)
        _ss_model.eval().to(DEVICE)

        ss_model = _ss_model
        print("Model loaded successfully.")
    except Exception as e:
        model_error = str(e)
        print(f"Error loading model: {traceback.format_exc()}")
    finally:
        model_loading = False


threading.Thread(target=load_model, daemon=True).start()


model_card = ModelCard(
    name="AudioSep",
    description="Separate any sound from a mixture using a natural language text description.",
    author="Xubo Liu, Qiuqiang Kong, Yan Zhao, Haohe Liu, Yi Yuan, Yuzhuo Liu, Rui Xia, Yuxuan Wang, Mark D. Plumbley, Wenwu Wang",
    tags=["audio separation", "text-queried", "source separation"],
)


@torch.inference_mode()
def process_fn(audio_path: str, text_query: str):
    if model_loading:
        raise gr.Error("Model is still loading, please wait a moment and try again.")
    if ss_model is None:
        raise gr.Error(f"Model failed to load: {model_error}")

    print(f"Separating [{audio_path}] with query [{text_query}]")

    mixture, _ = librosa.load(audio_path, sr=32000, mono=True)  # model expects 32k audio

    conditions = query_encoder.get_query_embed(
        modality="text",
        text=[text_query],
        device=DEVICE,
    )

    input_dict = {
        "mixture": torch.tensor(mixture, dtype=torch.float32)[None, None, :].to(DEVICE),
        "condition": conditions,
    }

    # Note: using ss_model.forward() directly.
    # chunk_inference() is not used — it has a latent self.sampling_rate AttributeError
    # in the original source and is only needed for very long audio to reduce peak memory.
    sep_segment = ss_model(input_dict)["waveform"]
    sep_np = sep_segment.squeeze(0).squeeze(0).cpu().numpy()

    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
        out_path = f.name
    sf.write(out_path, sep_np, 32000)

    return out_path


with gr.Blocks() as demo:
    input_components = [
        gr.Audio(type="filepath", label="Input Audio (Mixture)").harp_required(True),
        gr.Textbox(label="Text Query", placeholder="e.g. a dog barking"),
    ]
    output_components = [
        gr.Audio(type="filepath", label="Separated Audio").set_info(
            "Separated audio at 32 kHz matching the text description."
        ),
    ]

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

demo.queue().launch()