import gradio as gr
import os
import tempfile
import subprocess

def separate_audio(audio_path, stem_count):

    head, tail = os.path.split(audio_path)
    gradio_temp_path = head
    audio_filename = tail.split('.')[0]

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

    result = subprocess.run(command)

    outputs = []
    if stem_count == 2:
        accompaniment_path = f"{gradio_temp_path}/separated_audio/{audio_filename}/accompaniment.wav"
        vocals_path = f"{gradio_temp_path}/separated_audio/{audio_filename}/vocals.wav"
        outputs.append({'description': 'Accompaniment', 'path': accompaniment_path})
        outputs.append({'description': 'Vocals', 'path': vocals_path})
    elif stem_count == 4:
        vocals_path = f"{gradio_temp_path}/separated_audio/{audio_filename}/vocals.wav"
        drums_path = f"{gradio_temp_path}/separated_audio/{audio_filename}/drums.wav"
        bass_path = f"{gradio_temp_path}/separated_audio/{audio_filename}/bass.wav"
        other_path = f"{gradio_temp_path}/separated_audio/{audio_filename}/other.wav"
        outputs.extend([
            {'description': 'Vocals', 'path': vocals_path},
            {'description': 'Drums', 'path': drums_path},
            {'description': 'Bass', 'path': bass_path},
            {'description': 'Other', 'path': other_path},
        ])
    elif stem_count == 5:
        piano_path = f"{gradio_temp_path}/separated_audio/{audio_filename}/piano.wav"
        vocals_path = f"{gradio_temp_path}/separated_audio/{audio_filename}/vocals.wav"
        drums_path = f"{gradio_temp_path}/separated_audio/{audio_filename}/drums.wav"
        bass_path = f"{gradio_temp_path}/separated_audio/{audio_filename}/bass.wav"
        other_path = f"{gradio_temp_path}/separated_audio/{audio_filename}/other.wav"
        outputs.extend([
            {'description': 'Vocals', 'path': vocals_path},
            {'description': 'Piano', 'path': piano_path},
            {'description': 'Drums', 'path': drums_path},
            {'description': 'Bass', 'path': bass_path},
            {'description': 'Other', 'path': other_path},
        ])
    return outputs

def separate_audio_gradio(audio_file, stem_count):
    with tempfile.NamedTemporaryFile(delete=False) as temp_file:        
        temp_file.write(audio_file.read())
        temp_file_path = temp_file.name
        
    separate_audios = separate_audio(temp_file_path, stem_count)
    
    outputs = []
    for audio in separate_audios:
        outputs.append((audio['description'], gr.outputs.Audio(audio['path'])))
    
    return outputs

input_audio = gr.inputs.File(label="Upload Audio File", type=["wav", "mp3"])
input_stem_count = gr.inputs.Radio([2, 4, 5], label="Stem count")

gr.Interface(
    fn=separate_audio_gradio,
    inputs=[input_audio, input_stem_count],
    outputs="label",
    title="Audio Separator",
    description="Separate audio into vocals and accompaniment using Spleeter",
    capture_session=True
).launch()