"""
Stem Splitter — separa cualquier canción en 4 pistas (voz, batería, bajo, otros)
con HTDemucs v4. Gratis. Esta demo procesa los primeros 30 s (CPU gratuita de HF);
la versión completa (canción entera) está en adrianmoreno-dev.com.

Por Adrián Moreno · https://adrianmoreno-dev.com
"""
import os
import subprocess
import tempfile

import gradio as gr

MAX_SECONDS = 30
MODEL = "htdemucs"
STEMS = [("vocals", "🎤 Voz"), ("drums", "🥁 Batería"), ("bass", "🎸 Bajo"), ("other", "🎹 Otros")]


def separate(audio_path):
    if not audio_path:
        raise gr.Error("Sube una canción primero.")

    workdir = tempfile.mkdtemp()
    trimmed = os.path.join(workdir, "input.wav")

    # 1) Recorta a 30 s y normaliza a wav estéreo 44.1 kHz (ffmpeg lee mp3/wav/flac/ogg).
    subprocess.run(
        ["ffmpeg", "-y", "-i", audio_path, "-t", str(MAX_SECONDS),
         "-ac", "2", "-ar", "44100", trimmed],
        check=True, capture_output=True)

    # 2) Separa con HTDemucs (CLI estable).
    subprocess.run(
        ["python", "-m", "demucs", "-n", MODEL, "-o", workdir, trimmed],
        check=True, capture_output=True)

    out_dir = os.path.join(workdir, MODEL, "input")
    results = []
    for stem, _label in STEMS:
        p = os.path.join(out_dir, f"{stem}.wav")
        results.append(p if os.path.exists(p) else None)
    return results


with gr.Blocks(title="Stem Splitter", theme=gr.themes.Soft(primary_hue="teal")) as demo:
    gr.Markdown(
        "# 🎵 Stem Splitter — separa cualquier canción en 4 pistas\n"
        "Sube una canción y la IA (**HTDemucs v4**) la separa en **voz, batería, bajo y melodía**. "
        "Gratis y open source. *Esta demo procesa los primeros 30 segundos (CPU gratuita); puede tardar ~1 min.*\n\n"
        "👉 Canción completa y remezcla en tiempo real en "
        "[adrianmoreno-dev.com](https://adrianmoreno-dev.com/demo/stem-splitter?utm_source=huggingface)"
    )
    inp = gr.Audio(type="filepath", label="Sube una canción (mp3, wav…)")
    btn = gr.Button("🎚️ Separar en stems", variant="primary")
    with gr.Row():
        out_v = gr.Audio(label="🎤 Voz")
        out_d = gr.Audio(label="🥁 Batería")
    with gr.Row():
        out_b = gr.Audio(label="🎸 Bajo")
        out_o = gr.Audio(label="🎹 Otros")
    btn.click(separate, inp, [out_v, out_d, out_b, out_o])
    gr.Markdown("<sub>Open source de **Adrián Moreno** · [portfolio](https://adrianmoreno-dev.com?utm_source=huggingface) · modelo HTDemucs v4 (MIT)</sub>")

if __name__ == "__main__":
    demo.launch()