import os
import gradio as gr
from audio_separator.separator import Separator
import shutil
import threading
import time
import uuid

# Configuration
# Directory where files will be stored temporarily
UPLOAD_DIR = "output_storage"
MAX_FILE_AGE = 1800 

# Ensure the storage directory exists
os.makedirs(UPLOAD_DIR, exist_ok=True)

def cleanup_scheduler():
    """
    Background task that runs forever. 
    It scans the UPLOAD_DIR every minute and deletes folders older than MAX_FILE_AGE.
    """
    while True:
        try:
            current_time = time.time()
            # Iterate over all items in the storage directory
            for filename in os.listdir(UPLOAD_DIR):
                file_path = os.path.join(UPLOAD_DIR, filename)
                
                # Check if it's a directory (we create one folder per request)
                if os.path.isdir(file_path):
                    # Get the folder's creation/modification time
                    file_age = current_time - os.path.getmtime(file_path)
                    
                    if file_age > MAX_FILE_AGE:
                        print(f"🧹 Cleanup: Removing old folder {filename} ({int(file_age)}s old)")
                        shutil.rmtree(file_path)
                        
        except Exception as e:
            print(f"Error during cleanup: {e}")
        
        # Wait for 60 seconds before checking again
        time.sleep(60)

# Start the cleanup thread immediately (Daemon means it dies when the main app dies)
threading.Thread(target=cleanup_scheduler, daemon=True).start()


def separate_audio(input_file, model_name):
    if input_file is None:
        return (None, None, None, None, None, None)

    # Create a unique subfolder for this specific request inside our managed directory
    # Using UUID ensures no filename collisions between users
    session_id = str(uuid.uuid4())
    session_dir = os.path.join(UPLOAD_DIR, session_id)
    os.makedirs(session_dir, exist_ok=True)

    try:
        separator = Separator(
            output_dir=session_dir,
            output_format="MP3",
            output_bitrate="128k",
        )

        separator.load_model(model_filename=model_name)
        output_files = separator.separate(input_file)

        stems = ["Vocals", "Drums", "Bass", "Other", "Guitar", "Piano"]
        results = []

        for stem in stems:
            # Find the file that matches the stem name
            match = next(
                (os.path.join(session_dir, f) for f in output_files if stem in f),
                None,
            )
            results.append(match)

        # Note: We do NOT call a delete function here anymore. 
        # The background scheduler will handle it later.
        return tuple(results)

    except Exception as e:
        # If the actual separation fails, clean up this specific failed folder immediately
        if os.path.exists(session_dir):
            shutil.rmtree(session_dir)
        raise RuntimeError(f"Separation failed: {e}") from e


# ------------------ UI ------------------

with gr.Blocks() as demo:
    gr.Markdown(
        """
        # 🎵 AI 6-Stem Audio Separator
        Split your music into high-quality individual tracks using **Demucs 6s**.
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            input_audio = gr.Audio(
                label="Upload Audio File",
                type="filepath",
                interactive=True,
            )

            model_dropdown = gr.Dropdown(
                choices=["htdemucs_6s.yaml"],
                value="htdemucs_6s.yaml",
                label="Select Model",
            )

            separate_btn = gr.Button("🚀 Separate Stems", variant="primary")

        with gr.Column(scale=2):
            gr.Markdown("### 🔊 Separation Results")

            with gr.Row():
                output_vocals = gr.Audio(label="🎤 Vocals", type="filepath")
                output_drums = gr.Audio(label="🥁 Drums", type="filepath")

            with gr.Row():
                output_bass = gr.Audio(label="🎸 Bass", type="filepath")
                output_guitar = gr.Audio(label="🎸 Guitar", type="filepath")

            with gr.Row():
                output_piano = gr.Audio(label="🎹 Piano", type="filepath")
                output_other = gr.Audio(label="🎼 Other", type="filepath")

    separate_btn.click(
        fn=separate_audio,
        inputs=[input_audio, model_dropdown],
        outputs=[
            output_vocals,
            output_drums,
            output_bass,
            output_other,
            output_guitar,
            output_piano,
        ],
    )

    gr.Markdown(
        "⚠️ **Note:** Processing may take 1–2 minutes. Files are automatically cleaned up after 1 hour."
    )


if __name__ == "__main__":
    demo.launch(
        ssr_mode=False,
        theme=gr.themes.Soft(
            primary_hue="indigo",
            secondary_hue="blue",
        ),
    )