"""
Test speed/quality: Lead vs Backing vocals separation
Zakłada, że na wejściu dajesz JUŻ wyizolowany stem "vocals" (np. z Twojego htdemucs_6s),
dokładnie tak jak produkuje go worker.py. Testujemy tylko DRUGI etap pipeline'u.

Modele testowane (najlepsze wg community UVR / nomadkaraoke, stan na 2026):
- mel_band_roformer_karaoke_aufr33_viperx_sdr_10.1956.ckpt  (MDXC / Mel-Band Roformer)
- UVR_MDXNET_KARA_2.onnx                                     (MDX-Net, szybszy, starszy)

Wynik: lead_vocals.wav + backing_vocals.wav + czas przetwarzania + RTF (real-time factor)
"""

import os
import time
import shutil
import tempfile
from pathlib import Path

import gradio as gr
import torchaudio

from audio_separator.separator import Separator

MODEL_DIR = os.environ.get("MODEL_DIR", "/tmp/audio-separator-models")
os.makedirs(MODEL_DIR, exist_ok=True)

MODEL_CHOICES = {
    "Mel-Band Roformer Karaoke (aufr33/viperx) - najlepsza jakość": "mel_band_roformer_karaoke_aufr33_viperx_sdr_10.1956.ckpt",
    "UVR-MDX-NET Karaoke 2 - szybszy, starsza architektura": "UVR_MDXNET_KARA_2.onnx",
}

# Cache załadowanych modeli w pamięci procesu, żeby nie przeładowywać za każdym razem
_loaded_separators = {}


def get_separator(model_filename: str, output_dir: str) -> Separator:
    if model_filename not in _loaded_separators:
        sep = Separator(
            output_dir=output_dir,
            model_file_dir=MODEL_DIR,
        )
        sep.load_model(model_filename=model_filename)
        _loaded_separators[model_filename] = sep
    else:
        # ten sam obiekt, ale trzeba zaktualizować output_dir na nowy job
        _loaded_separators[model_filename].output_dir = output_dir
    return _loaded_separators[model_filename]


def get_audio_duration(path: str) -> float:
    info = torchaudio.info(path)
    return info.num_frames / info.sample_rate


def run_separation(vocals_file, model_label, progress=gr.Progress()):
    if vocals_file is None:
        raise gr.Error("Wgraj plik ze stemem wokalu (np. vocals.wav z Twojego htdemucsa).")

    model_filename = MODEL_CHOICES[model_label]

    work_dir = tempfile.mkdtemp(prefix="leadback_")
    try:
        duration_sec = get_audio_duration(vocals_file)

        progress(0.05, desc=f"Ładowanie modelu {model_filename} (pierwsze uruchomienie może pobierać wagi)...")
        t_load_start = time.time()
        separator = get_separator(model_filename, work_dir)
        t_load = time.time() - t_load_start

        progress(0.3, desc="Separacja lead / backing w toku...")
        t_sep_start = time.time()
        output_files = separator.separate(vocals_file)
        t_sep = time.time() - t_sep_start

        progress(0.95, desc="Finalizacja...")

        # audio-separator zwraca listę plików wyjściowych - rozróżniamy po nazwie
        lead_path, backing_path = None, None
        for f in output_files:
            full_path = f if os.path.isabs(f) else os.path.join(work_dir, f)
            lower = full_path.lower()
            if "instrumental" in lower or "backing" in lower or "no_vocal" in lower:
                backing_path = full_path
            else:
                lead_path = full_path

        # fallback, gdyby nazwy się nie zgadzały - bierzemy po kolei
        if (lead_path is None or backing_path is None) and len(output_files) >= 2:
            paths = [f if os.path.isabs(f) else os.path.join(work_dir, f) for f in output_files]
            lead_path, backing_path = paths[0], paths[1]

        rtf = t_sep / duration_sec if duration_sec > 0 else 0

        report = (
            f"**Model:** {model_filename}\n\n"
            f"**Długość utworu:** {duration_sec:.1f}s\n\n"
            f"**Czas ładowania modelu:** {t_load:.1f}s (tylko przy pierwszym uruchomieniu / zmianie modelu)\n\n"
            f"**Czas samej separacji:** {t_sep:.1f}s\n\n"
            f"**RTF (real-time factor):** {rtf:.2f}x "
            f"({'szybciej niż realtime' if rtf < 1 else 'wolniej niż realtime'})\n\n"
            f"_Wskazówka: RTF x2 znaczy, że 3-minutowy utwór = ~6s przetwarzania na tym sprzęcie._"
        )

        return lead_path, backing_path, report

    except Exception as e:
        raise gr.Error(f"Błąd separacji: {e}")


with gr.Blocks(title="Test: Lead vs Backing Vocals") as demo:
    gr.Markdown(
        "# 🎤 Test separacji Lead / Backing Vocals\n"
        "Wgraj **już wyizolowany stem wokalu** (dokładnie taki jak produkuje Twój `worker.py` - plik `vocals.wav`/`.mp3`). "
        "Ta appka testuje TYLKO drugi etap pipeline'u (lead vs backing), żebyś zobaczył realny narzut czasowy "
        "przed wdrożeniem na produkcję.\n\n"
        "⚠️ Pierwsze uruchomienie danego modelu pobierze jego wagi (setki MB) - kolejne testy będą szybsze."
    )

    with gr.Row():
        with gr.Column():
            vocals_input = gr.Audio(label="Stem wokalu (wejście)", type="filepath")
            model_dropdown = gr.Dropdown(
                choices=list(MODEL_CHOICES.keys()),
                value=list(MODEL_CHOICES.keys())[0],
                label="Model do testu",
            )
            run_btn = gr.Button("▶️ Uruchom separację", variant="primary")

        with gr.Column():
            lead_output = gr.Audio(label="Lead vocal (wynik)")
            backing_output = gr.Audio(label="Backing vocals (wynik)")
            report_output = gr.Markdown()

    run_btn.click(
        fn=run_separation,
        inputs=[vocals_input, model_dropdown],
        outputs=[lead_output, backing_output, report_output],
    )

if __name__ == "__main__":
    demo.queue().launch(
        server_name="0.0.0.0",
        server_port=int(os.environ.get("PORT", 7860)),
        ssr_mode=False,
    )