import os
import gradio as gr
from audio_separator.separator import Separator

# Define models
models = {
    "Roformer": {
        'BS-Roformer-Viperx-1297.ckpt': 'model_bs_roformer_ep_317_sdr_12.9755.ckpt',
        'BS-Roformer-Viperx-1296.ckpt': 'model_bs_roformer_ep_368_sdr_12.9628.ckpt',
        'BS-Roformer-Viperx-1053.ckpt': 'model_bs_roformer_ep_937_sdr_10.5309.ckpt',
        'Mel-Roformer-Viperx-1143.ckpt': 'model_mel_band_roformer_ep_3005_sdr_11.4360.ckpt'
    },
    "MDX23C": [
        'MDX23C_D1581.ckpt',
        'MDX23C-8KFFT-InstVoc_HQ.ckpt',
        'MDX23C-8KFFT-InstVoc_HQ_2.ckpt',
    ],
    "MDXNet": [
        'UVR-MDX-NET-Inst_full_292.onnx',
        'UVR-MDX-NET_Inst_187_beta.onnx',
        # Add remaining models as needed
    ],
    "VRArch": [
        '1_HP-UVR.pth',
        '2_HP-UVR.pth',
        # Add remaining models as needed
    ],
    "Demucs": [
        'htdemucs_ft.yaml',
        'htdemucs.yaml',
        # Add remaining models as needed
    ]
}

# Function for separating the audio
def separate_audio(audio, model_name):
    if audio is None:
        return None, None

    # Get the selected model path
    model_path = None
    for model_group in models.values():
        if isinstance(model_group, dict):
            model_path = model_group.get(model_name)
            if model_path:
                break
        elif model_name in model_group:
            model_path = model_name
            break

    if model_path is None:
        return None, None

    # Set up output directory and models for separation
    output_dir = "./output"
    os.makedirs(output_dir, exist_ok=True)

    separator = Separator(output_dir=output_dir)

    # Define output paths
    vocals_path = os.path.join(output_dir, 'Vocals.wav')
    instrumental_path = os.path.join(output_dir, 'Instrumental.wav')

    with gr.Progress() as progress:
        progress(0, "Loading model...")
        separator.load_model(model_filename=model_path)
        progress(20, "Model loaded. Starting separation...")

        # Step 3: Splitting track into Vocal and Instrumental
        voc_inst = separator.separate(audio)

        # Check if separation was successful
        if len(voc_inst) != 2:
            return None, None
        
        # Save the separated files
        os.rename(voc_inst[0], instrumental_path)
        os.rename(voc_inst[1], vocals_path)

        progress(100, "Separation complete!")
    
    # Return paths to the processed files
    return instrumental_path, vocals_path

# Define the Gradio Interface
with gr.Blocks(theme="NoCrypt/miku@1.2.2") as demo:
    gr.Markdown("# Audio Separator Gradio Demo")

    with gr.Row():
        with gr.Column():
            link_input = gr.Audio(label="Upload Audio File", type="filepath")
            model_dropdown = gr.Dropdown(label="Select Model", choices=list(
                models["Roformer"].keys()) + models["MDX23C"] + models["MDXNet"] + models["VRArch"] + models["Demucs"])
            separate_button = gr.Button("Separate Audio")
        
        with gr.Column():
            instrumental_output = gr.Audio(label="Instrumental Output", type="filepath")
            vocals_output = gr.Audio(label="Vocals Output", type="filepath")
            
    # Define button functionality
    separate_button.click(
        separate_audio,
        inputs=[link_input, model_dropdown],
        outputs=[
            instrumental_output,
            vocals_output,
        ]
    )

# Launch the Gradio app
demo.launch(debug=True, share=True)