import spaces  # imported first, as ZeroGPU requires

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

import os
import tempfile
import traceback
from pathlib import Path

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, SampleRate, Channels
from pyharp.tags import Subcategory

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

model_card = ModelCard(
    name="AudioSep",
    description=(
        "Separate any sound from a mixture using a natural language text "
        "description. Separation runs at 32 kHz in mono. The result is "
        "resampled back to the input's rate and copied across its channels, so "
        "the output has the format it came in with but not its stereo image. "
        "Repository: https://github.com/Audio-AGI/AudioSep Paper: "
        "https://arxiv.org/abs/2308.05037"
    ),
    author="Xubo Liu, Qiuqiang Kong, Yan Zhao, Haohe Liu, Yi Yuan, Yuzhuo Liu, Rui Xia, Yuxuan Wang, Mark D. Plumbley, Wenwu Wang",
    tags=[Subcategory.TARGET_SOURCE_EXTRACTION, SampleRate(32000), Channels(1), 'text-queried', 'domain:general-audio'],
)

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

# A lossless input keeps its container. A lossy one comes back as WAV, since encoding the result
# lossily again would expose the artifacts of the first encode and add a second set.
LOSSLESS = {".wav", ".flac"}

ss_model = None
query_encoder = None
model_error = None


def load_model():
    global ss_model, query_encoder, 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",
            revision="0807a713283035d8b0d0628f1b8e3d6580459f7b",
        )
        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",
            revision="0807a713283035d8b0d0628f1b8e3d6580459f7b",
        )

        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()}")


# pyharp runs process_fn in a worker, which imports this file as __mp_main__ and starts with the
# server, so the model is loaded there alone, before process_fn can reserve a GPU. A failure is
# kept rather than raised, so that each request can report it.
if __name__ != "__main__":
    load_model()


def process_fn(audio_path: str, text_query: str):
    # Checked before the GPU is reserved, so a request that cannot run spends none of the caller's
    # ZeroGPU allowance
    if ss_model is None:
        raise gr.Error(f"Model failed to load: {model_error}")
    if not audio_path:
        raise gr.Error("Provide the mixture to separate.")
    try:
        info = sf.info(audio_path)
    except Exception as exc:
        raise gr.Error(f"Could not read the audio file: {exc}") from exc
    if info.frames == 0:
        raise gr.Error("The audio file is empty.")
    if not (text_query or "").strip():
        raise gr.Error("Describe the sound to separate, such as \"acoustic guitar\".")

    return separate(audio_path, text_query.strip())


@spaces.GPU(duration=120)
@torch.inference_mode()
def separate(audio_path: str, text_query: str):

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

    # The model expects 32 kHz mono. Both are restored after inference, so the
    # result drops into the session it came from rather than at the model's rate
    info = sf.info(audio_path)
    original_sr = info.samplerate
    original_channels = info.channels

    mixture, _ = librosa.load(audio_path, sr=32000, mono=True)

    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()

    if original_sr != 32000:
        sep_np = librosa.resample(sep_np, orig_sr=32000, target_sr=original_sr)

    # Separation here is mono, so the channels are rebuilt by copying rather
    # than recovered, as the model card says
    if original_channels > 1:
        sep_np = np.tile(sep_np[:, None], (1, original_channels))

    input_ext = Path(audio_path).suffix.lower()
    with tempfile.NamedTemporaryFile(suffix=input_ext if input_ext in LOSSLESS else ".wav",
                                     delete=False) as f:
        out_path = f.name
    sf.write(out_path, sep_np, original_sr)

    return out_path


if __name__ == "__main__":
    with gr.Blocks() as demo:
        input_components = [
            gr.Audio(type="filepath", label="Input Audio (Mixture)").harp_required(True).set_info((
                "The mixture containing the sound to isolate."
            )),
            gr.Textbox(label="Text Query", placeholder="e.g. a dog barking",
            info=(
                "Describe the sound to isolate, such as \"acoustic guitar\" or \"a "
                "dog barking\". Plain words work better than studio terms."
            )),
        ]
        output_components = [
            gr.Audio(type="filepath", label="Separated Audio").set_info(
                "The sound matching the description, at the sample rate and channel count of "
                "the input."
            ),
        ]

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

    demo.queue().launch(pwa=True, show_error=True, ssr_mode=False)