# app.py
import os
import shutil
import subprocess
import uuid
import gradio as gr

OUTPUT_DIR = "outputs"

os.makedirs(OUTPUT_DIR, exist_ok=True)


def separate_vocals(audio_file):
    try:
        if audio_file is None:
            return None

        uid = str(uuid.uuid4())

        work_dir = os.path.join(OUTPUT_DIR, uid)
        os.makedirs(work_dir, exist_ok=True)

        input_audio = os.path.join(work_dir, "input.wav")

        # Copy uploaded file
        shutil.copy(audio_file, input_audio)

        # Run Demucs
        command = [
            "python",
            "-m",
            "demucs.separate",
            "-n",
            "htdemucs",
            "--two-stems=vocals",
            input_audio,
            "-o",
            work_dir
        ]

        subprocess.run(command, check=True)

        # Find vocal file
        vocals_path = None

        demucs_folder = os.path.join(
            work_dir,
            "htdemucs",
            "input"
        )

        for file in os.listdir(demucs_folder):
            if "vocals" in file:
                vocals_path = os.path.join(demucs_folder, file)
                break

        return vocals_path

    except Exception as e:
        return str(e)


app = gr.Interface(
    fn=separate_vocals,
    inputs=gr.Audio(
        type="filepath",
        label="Upload Audio"
    ),
    outputs=gr.Audio(
        label="Vocals Only Output"
    ),
    title="AI Vocal Remover",
    description="Upload audio and remove background music/instruments."
)

app.launch(server_name="0.0.0.0", server_port=7860)