import gradio as gr
import os
import tempfile
import subprocess

def separate_audio(audio_path, stem_count):
    print(f"{audio_path=}")
    head, tail = os.path.split(audio_path)

    gradio_temp_path = head
    audio_filename = tail.split('.')[0]
    print(f"{gradio_temp_path=}")
    print(f"{audio_filename=}")
    print(f"{stem_count=}")

    command = f"spleeter separate -p spleeter:{stem_count}stems {audio_path}"
    command = command.split()
    print(f"{command=}")

    result = subprocess.run(command)
    print(result)

    paths = []
    if stem_count == 2:
        paths.append(('Accompaniment', f"{gradio_temp_path}/separated_audio/{audio_filename}/accompaniment.wav"))
        paths.append(('Vocals', f"{gradio_temp_path}/separated_audio/{audio_filename}/vocals.wav"))
    elif stem_count == 4 or stem_count == 5:
        paths.append(('Vocals', f"{gradio_temp_path}/separated_audio/{audio_filename}/vocals.wav"))
        paths.append(('Drums', f"{gradio_temp_path}/separated_audio/{audio_filename}/drums.wav"))
        paths.append(('Bass', f"{gradio_temp_path}/separated_audio/{audio_filename}/bass.wav"))
        paths.append(('Other', f"{gradio_temp_path}/separated_audio/{audio_filename}/other.wav"))
        if stem_count == 5:
            paths.append(('Piano', f"{gradio_temp_path}/separated_audio/{audio_filename}/piano.wav"))

    return paths

def gradio_interface(audio_file, stem_count):
    with tempfile.NamedTemporaryFile(delete=False) as temp_file:
        temp_file.write(audio_file)
        temp_file_path = temp_file.name

    separated_audios = separate_audio(temp_file_path, stem_count)
    return [gr.Audio(file_path, label=description) for description, file_path in separated_audios]

iface = gr.Interface(
    fn=gradio_interface,
    inputs=[gr.Audio(type="filepath"), gr.Radio([2, 4, 5])],
    outputs=[gr.Audio(label="Output") for _ in range(5)],
)

iface.launch()