import gradio as gr
import subprocess
import os
import shutil

def separate_audio(input_file):
    if input_file is None:
        return None, None, "Please upload a file first."
    
    # 1. Setup paths
    output_dir = "separated"
    if os.path.exists(output_dir):
        shutil.rmtree(output_dir) # Clean up previous run
    os.makedirs(output_dir, exist_ok=True)
    
    print(f"Processing: {input_file}")
    
    # 2. Run Demucs (Command Line Interface inside Python)
    # We use "htdemucs" (high quality) and 2 stems (vocals/accompaniment)
    command = [
        "demucs",
        "--two-stems=vocals",
        "-n", "htdemucs",
        "--out", output_dir,
        input_file
    ]
    
    try:
        subprocess.run(command, check=True)
    except subprocess.CalledProcessError as e:
        return None, None, f"Error during separation: {str(e)}"

    # 3. Locate Output Files
    # Demucs structure: separated/htdemucs/{filename}/vocals.wav
    filename_no_ext = os.path.splitext(os.path.basename(input_file))[0]
    target_folder = os.path.join(output_dir, "htdemucs", filename_no_ext)
    
    vocals_path = os.path.join(target_folder, "vocals.wav")
    bg_path = os.path.join(target_folder, "no_vocals.wav")
    
    if os.path.exists(vocals_path) and os.path.exists(bg_path):
        return vocals_path, bg_path, "✅ Separation Complete!"
    else:
        return None, None, "❌ Could not find output files."

# --- The Interface ---
with gr.Blocks(title="Indic Dubbing Studio") as demo:
    gr.Markdown("# 🎙️ Indic Dubbing Studio - Phase 1")
    gr.Markdown("Step 1: Separate your video into **Vocals** (for translation) and **Background** (for mixing).")
    
    with gr.Row():
        with gr.Column():
            input_file = gr.Audio(type="filepath", label="Upload Video or Audio", sources=["upload"])
            process_btn = gr.Button("Split Audio (Demucs)", variant="primary")
        
        with gr.Column():
            status_output = gr.Label(label="Status")
            out_vocals = gr.Audio(label="Vocals (Clean Speech)")
            out_bg = gr.Audio(label="Background (Music/SFX)")

    process_btn.click(
        fn=separate_audio,
        inputs=[input_file],
        outputs=[out_vocals, out_bg, status_output]
    )

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