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

import traceback
import contextlib
import tempfile
import torch
import os
import threading
# import spaces  # uncomment when running on ZeroGPU
import librosa
import soundfile as sf
import gradio as gr
from huggingface_hub import snapshot_download
from pyharp import ModelCard, build_endpoint

from anyaccomp.inference_utils import Sing2SongInferencePipeline

repo_id = "amphion/anyaccomp"
base_dir = os.path.dirname(os.path.abspath(__file__))
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# Model loading state. Populated by the background thread.
inference_pipeline = None
model_loading = True
model_error = None


def load_model():
    global inference_pipeline, model_loading, model_error
    try:
        checkpoint_marker = os.path.join(base_dir, "pretrained", "flow_matching")
        if not os.path.exists(checkpoint_marker):
            print(f"Downloading model files from {repo_id}...", flush=True)
            model_dir = snapshot_download(repo_id=repo_id, local_dir=base_dir)
            print(f"Model files downloaded to: {model_dir}", flush=True)
        else:
            model_dir = base_dir
            print("Model files already present, skipping download.", flush=True)

        cfg_path = os.path.join(base_dir, "config/flow_matching.json")
        vocoder_cfg_path = os.path.join(base_dir, "config/vocoder.json")
        checkpoint_path = os.path.join(model_dir, "pretrained/flow_matching")
        vocoder_checkpoint_path = os.path.join(model_dir, "pretrained/vocoder")

        print("Initializing AnyAccomp InferencePipeline...", flush=True)
        pipeline = Sing2SongInferencePipeline(
            checkpoint_path, cfg_path,
            vocoder_checkpoint_path, vocoder_cfg_path,
            device=DEVICE,
        )
        pipeline.sample_rate = 24000
        inference_pipeline = pipeline
        print("Model loaded successfully.", flush=True)
    except Exception as e:
        model_error = str(e)
        print(f"Error loading model: {e}", flush=True)
    finally:
        model_loading = False


# Create a ModelCard describing this endpoint for HARP.
model_card = ModelCard(
    name="AnyAccomp",
    description="Upload a 3-30 second vocal or instrument track and the model will generate an accompaniment.",
    author="Amphion",
    tags=["audio", "music", "accompaniment"],
)

# Start model loading in the background so the server can start immediately.
threading.Thread(target=load_model, daemon=True).start()


# @spaces.GPU  # uncomment when running on ZeroGPU
@torch.inference_mode()
def process_fn(vocal_filepath, n_timesteps):  # add , cfg_scale to re-enable
    if model_loading:
        raise gr.Error("Model is still loading, please wait a moment and try again.")
    if inference_pipeline is None:
        raise gr.Error(f"Model failed to load: {model_error}")

    if vocal_filepath is None:
        raise gr.Error("Please upload a vocal audio file.")

    try:
        duration = librosa.get_duration(path=vocal_filepath)
        if not (3 <= duration <= 30):
            raise gr.Error("Audio duration must be between 3 and 30 seconds.")
    except Exception as e:
        raise gr.Error(f"Cannot read audio file or get duration: {e}")

    try:
        vocal_audio, _ = librosa.load(vocal_filepath, sr=24000, mono=True)
        vocal_tensor = torch.tensor(vocal_audio).unsqueeze(0).to(DEVICE)

        vocal_mel = inference_pipeline.encode_vocal(vocal_tensor)

        autocast_ctx = torch.amp.autocast("cuda", dtype=torch.bfloat16) if torch.cuda.is_available() else contextlib.nullcontext()
        with autocast_ctx:
            mel = inference_pipeline.model.reverse_diffusion(
                vocal_mel=vocal_mel,
                n_timesteps=int(n_timesteps),
                cfg=3.0,  # replace with cfg_scale to re-enable
            )

        mel = mel.float()
        wav = inference_pipeline._generate_audio(mel)
        wav = wav.squeeze().detach().cpu().numpy()

        wav = librosa.util.fix_length(data=wav, size=len(vocal_audio))
        mixture_wav = wav + vocal_audio

        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
            accompaniment_path = f.name
        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
            mixture_path = f.name

        sf.write(accompaniment_path, wav, 24000)
        sf.write(mixture_path, mixture_wav, 24000)

        return accompaniment_path, mixture_path

    except Exception as e:
        traceback.print_exc()
        raise gr.Error(f"An error occurred during processing: {e}")


with gr.Blocks() as demo:
    input_components = [
        gr.Audio(
            type="filepath",
            label="Input Vocal or Instrument Audio",
        ).harp_required(True),
        gr.Slider(
            minimum=10, maximum=100, value=50, step=1,
            label="Inference Steps (n_timesteps)",
            info="Number of diffusion steps. More steps = slower but higher quality",
        ),
        # gr.Slider(minimum=1.0, maximum=10.0, value=3.0, step=0.1, label="CFG Scale", info="Guidance strength — higher values follow the input melody more closely."),
    ]

    output_components = [
        gr.Audio(type="filepath", label="Generated Accompaniment").set_info(
            "The generated instrumental accompaniment."
        ),
        gr.Audio(type="filepath", label="Mixture (Vocal + Accompaniment)").set_info(
            "Vocal mixed with the generated accompaniment."
        ),
    ]

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


demo.queue().launch()