import os
import gradio as gr
import subprocess
import shutil

def separate_vocal(audio):
    if audio is None:
        return None, None

    # 1. Fanadiovana ny toerana hiasana
    output_dir = "results"
    if os.path.exists(output_dir):
        shutil.rmtree(output_dir)
    
    try:
        # 2. Ampiasaina ny model 'htdemucs_light' (maivana kokoa ho an'ny CPU maimaim-poana)
        # Nampiana --jobs 1 mba tsy hikatso ny serveur
        print(f"Fikarakarana ny hira: {audio}")
        
        # Mampiasa subprocess mba handehanan'ny baiko Demucs
        cmd = [
            "demucs", 
            "--two-stems", "vocals", 
            "-n", "htdemucs_light", 
            "-o", output_dir, 
            "--jobs", "1",
            audio
        ]
        
        subprocess.run(cmd, check=True)

        # 3. Mitady ny file nivoaka
        # Mitady any anatin'ny results/htdemucs_light/...
        model_folder = os.path.join(output_dir, "htdemucs_light")
        song_folder_name = os.listdir(model_folder)[0]
        song_folder_path = os.path.join(model_folder, song_folder_name)
        
        vocal_path = os.path.join(song_folder_path, "vocals.wav")
        instr_path = os.path.join(song_folder_path, "no_vocals.wav")

        return vocal_path, instr_path

    except Exception as e:
        print(f"Hadisoana nipoitra: {e}")
        return None, None

# Interface tsotra (Nesorina ilay 'allow_flagging' niteraka olana)
demo = gr.Interface(
    fn=separate_vocal,
    inputs=gr.Audio(type="filepath", label="Ampidiro ny hira"),
    outputs=[
        gr.Audio(label="Feon'olona (Vocals)"),
        gr.Audio(label="Mozika (Instrumental)")
    ],
    title="AI Vocal Separator",
    description="Andraso kely rehefa avy manindry 'Submit'. Ny hira lava dia mety mila minitra vitsy."
)

if __name__ == "__main__":
    demo.launch()