import os
import tempfile
import base64
import gradio as gr
import demucs.separate
import librosa

def separate_audio(file):
    """
    Accepts an audio file, runs Demucs, returns instrumental and vocals as base64.
    """
    if file is None:
        raise gr.Error("No audio file provided")

    with tempfile.TemporaryDirectory() as tmpdir:
        # Save uploaded file
        input_path = os.path.join(tmpdir, "input." + file.name.split(".")[-1])
        with open(input_path, "wb") as f:
            f.write(file.read())

        output_dir = os.path.join(tmpdir, "output")
        os.makedirs(output_dir, exist_ok=True)

        try:
            demucs.separate.main(
                ["--model", "htdemucs_ft", "--out", output_dir, input_path]
            )
        except Exception as e:
            raise gr.Error(f"Demucs failed: {str(e)}")

        base_name = os.path.splitext(os.path.basename(input_path))[0]
        sep_dir = os.path.join(output_dir, "htdemucs_ft", base_name)

        vocals_path = os.path.join(sep_dir, "vocals.wav")
        instrumental_path = os.path.join(sep_dir, "no_vocals.wav")

        if not os.path.exists(vocals_path) or not os.path.exists(instrumental_path):
            raise gr.Error("Separation output missing")

        def encode_wav(path):
            with open(path, "rb") as f:
                return base64.b64encode(f.read()).decode("utf-8")

        def get_info(path):
            y, sr = librosa.load(path, sr=None)
            return {"duration": len(y) / sr, "sample_rate": sr}

        return {
            "instrumental_b64": encode_wav(instrumental_path),
            "vocals_b64": encode_wav(vocals_path),
            "instrumental_info": get_info(instrumental_path),
            "vocals_info": get_info(vocals_path),
        }

# Gradio interface
iface = gr.Interface(
    fn=separate_audio,
    inputs=gr.Audio(type="file", label="Upload your song"),
    outputs=gr.JSON(label="Result"),
    title="Kaio Vocal Separator",
    description="Upload a song, get instrumental and vocals (base64).",
    api_name="separate"   # endpoint: /api/predict/separate
)

iface.launch()