import gradio as gr
import os
import subprocess
import spaces
from typing import Tuple, List, Dict
from pydub import AudioSegment
from rich.console import Console
from rich.panel import Panel
from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TimeRemainingColumn
from rich.text import Text
import time

console = Console()

@spaces.GPU
def inference(audio_file: str, model_name: str, vocals: bool, drums: bool, bass: bool, other: bool, mp3: bool, mp3_bitrate: int) -> Tuple[str, List[str], gr.HTML]:
    log_messages = []

    def stream_log(message, style=""):
        formatted_message = f"[{model_name}] {message}"
        log_messages.append(formatted_message)
        return gr.HTML(f"<pre style='margin-bottom: 0;{style}'>{formatted_message}</pre>")

    yield None, None, stream_log("Initializing Demucs...", "color: #4CAF50; font-weight: bold;")
    time.sleep(1)  # Simulate initialization time

    yield None, None, stream_log("Loading audio file...", "color: #2196F3;")
    time.sleep(0.5)  # Simulate loading time
    
    if audio_file is None:
        yield None, None, stream_log("Error: No audio file provided", "color: #F44336;")
        raise gr.Error("Please upload an audio file")

    # Use absolute paths
    base_output_dir = os.path.abspath("separated")
    output_dir = os.path.join(base_output_dir, model_name, os.path.splitext(os.path.basename(audio_file))[0])
    os.makedirs(output_dir, exist_ok=True)

    # Construct the Demucs command with full paths
    cmd = [
        "python", "-m", "demucs", 
        "--out", base_output_dir,
        "-n", model_name,
        audio_file
    ]

    yield None, None, stream_log("Preparing separation process...", "color: #FF9800;")
    time.sleep(0.5)  # Simulate preparation time

    try:
        # Run the Demucs command
        process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True)
        
        # Simulate a loading animation
        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            BarColumn(),
            TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
            TimeRemainingColumn(),
            console=console
        ) as progress:
            task = progress.add_task("[cyan]Separating stems...", total=100)
            while process.poll() is None:
                progress.update(task, advance=1)
                time.sleep(0.1)
            progress.update(task, completed=100)

        if process.returncode != 0:
            error_output = process.stderr.read()
            yield None, None, stream_log(f"Error: Separation failed", "color: #F44336;")
            raise gr.Error(f"Demucs separation failed. Check the logs for details.")

    except Exception as e:
        yield None, None, stream_log(f"Unexpected error: {str(e)}", "color: #F44336;")
        raise gr.Error(f"An unexpected error occurred: {str(e)}")

    yield None, None, stream_log("Separation completed successfully!", "color: #4CAF50; font-weight: bold;")
    time.sleep(0.5)  # Pause for effect

    yield None, None, stream_log("Processing stems...", "color: #9C27B0;")
    time.sleep(0.5)  # Simulate processing time

    # Change the stem search directory using full path
    stem_search_dir = os.path.join(base_output_dir, model_name, os.path.splitext(os.path.basename(audio_file))[0])
    yield None, None, stream_log(f"Searching for stems in: {stem_search_dir}")

    stems: Dict[str, str] = {}
    for stem in ["vocals", "drums", "bass", "other"]:
        stem_path = os.path.join(stem_search_dir, f"{stem}.wav")
        yield None, None, stream_log(f"Checking for {stem} stem at: {stem_path}")
        if os.path.exists(stem_path):
            stems[stem] = stem_path
            yield None, None, stream_log(f"Found {stem} stem")
        else:
            yield None, None, stream_log(f"Warning: {stem} stem not found")

    if not stems:
        yield None, None, stream_log("Error: No stems found. Checking alternative directory...")
        stem_search_dir = os.path.join(base_output_dir, model_name)
        for stem in ["vocals", "drums", "bass", "other"]:
            stem_path = os.path.join(stem_search_dir, f"{stem}.wav")
            yield None, None, stream_log(f"Checking for {stem} stem at: {stem_path}")
            if os.path.exists(stem_path):
                stems[stem] = stem_path
                yield None, None, stream_log(f"Found {stem} stem")
            else:
                yield None, None, stream_log(f"Warning: {stem} stem not found")

    yield None, None, stream_log(f"All found stems: {list(stems.keys())}")

    selected_stems: List[str] = []
    for stem, selected in zip(["vocals", "drums", "bass", "other"], [vocals, drums, bass, other]):
        if selected:
            yield None, None, stream_log(f"{stem} is selected by user")
            if stem in stems:
                selected_stems.append(stems[stem])
                yield None, None, stream_log(f"Selected {stem} stem for mixing")
            else:
                yield None, None, stream_log(f"Warning: {stem} was selected but not found")

    yield None, None, stream_log(f"Final selected stems: {selected_stems}")

    if not selected_stems:
        yield None, None, stream_log("Error: No stems selected for mixing", "color: #F44336;")
        raise gr.Error("Please select at least one stem to mix and ensure it was successfully separated.")

    output_file: str = os.path.join(output_dir, "mixed.wav")
    yield None, None, stream_log("Mixing selected stems...", "color: #FF5722;")
    time.sleep(0.5)  # Simulate mixing time
    
    mixed_audio: AudioSegment = AudioSegment.empty()
    for stem_path in selected_stems:
        mixed_audio += AudioSegment.from_wav(stem_path)
    mixed_audio.export(output_file, format="wav")

    if mp3:
        yield None, None, stream_log(f"Converting to MP3...", "color: #795548;")
        time.sleep(0.5)  # Simulate conversion time
        mp3_output_file: str = os.path.splitext(output_file)[0] + ".mp3"
        mixed_audio.export(mp3_output_file, format="mp3", bitrate=str(mp3_bitrate) + "k")
        output_file = mp3_output_file

    yield None, None, stream_log("Process completed successfully!", "color: #4CAF50; font-weight: bold;")
    yield output_file, list(stems.values()), gr.HTML(
        Panel.fit(
            Text("Separation and mixing completed successfully!", style="bold green"),
            title="Demucs Result",
            border_style="green"
        ).render()
    )

# Define the Gradio interface
with gr.Blocks() as iface:
    gr.Markdown("# Demucs Music Source Separation and Mixing")
    gr.Markdown("Separate vocals, drums, bass, and other instruments from your music using Demucs and mix the selected stems.")
    
    with gr.Row():
        with gr.Column(scale=1):
            audio_input = gr.Audio(type="filepath", label="Upload Audio File")
            model_dropdown = gr.Dropdown(
                ["htdemucs", "htdemucs_ft", "htdemucs_6s", "hdemucs_mmi", "mdx", "mdx_extra", "mdx_q", "mdx_extra_q"],
                label="Model Name",
                value="htdemucs_ft"
            )
            with gr.Row():
                vocals_checkbox = gr.Checkbox(label="Vocals", value=True)
                drums_checkbox = gr.Checkbox(label="Drums", value=True)
            with gr.Row():
                bass_checkbox = gr.Checkbox(label="Bass", value=True)
                other_checkbox = gr.Checkbox(label="Other", value=True)
            mp3_checkbox = gr.Checkbox(label="Save as MP3", value=False)
            mp3_bitrate = gr.Slider(128, 320, step=32, label="MP3 Bitrate", visible=False)
            submit_btn = gr.Button("Process", variant="primary")

        with gr.Column(scale=1):
            output_audio = gr.Audio(type="filepath", label="Processed Audio (Mixed)")
            stems_output = gr.File(label="Individual Stems", file_count="multiple")
            separation_log = gr.HTML()

    submit_btn.click(
        fn=inference,
        inputs=[audio_input, model_dropdown, vocals_checkbox, drums_checkbox, bass_checkbox, other_checkbox, mp3_checkbox, mp3_bitrate],
        outputs=[output_audio, stems_output, separation_log]
    )
    
    mp3_checkbox.change(
        fn=lambda mp3: gr.update(visible=mp3),
        inputs=mp3_checkbox,
        outputs=mp3_bitrate
    )

# Launch the Gradio interface
iface.launch()