import gradio as gr
import subprocess
import os
import shutil
from pydub import AudioSegment
import uuid
import time

# --- CSS EXACT MATCH FROM YOUR STYLE.CSS ---
custom_css = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700;800&display=swap');

body, .gradio-container {
    background-color: #0f0f23 !important;
    font-family: "Inter", sans-serif !important;
    color: #ffffff !important;
}

/* Background Gradient Fix */
div.gradio-container {
    background: radial-gradient(circle at 20% 30%, rgba(102, 126, 234, 0.15) 0%, transparent 50%),
                radial-gradient(circle at 80% 70%, rgba(118, 75, 162, 0.15) 0%, transparent 50%),
                radial-gradient(circle at 50% 50%, rgba(245, 87, 108, 0.1) 0%, transparent 70%) !important;
    background-color: #0f0f23 !important;
}

/* Glass Card Styling */
.glass-card {
    background: rgba(255, 255, 255, 0.05) !important;
    backdrop-filter: blur(20px) !important;
    border: 1px solid rgba(255, 255, 255, 0.1) !important;
    border-radius: 24px !important;
    padding: 40px !important;
    box-shadow: 0 8px 32px rgba(0, 0, 0, 0.3), inset 0 1px 0 rgba(255, 255, 255, 0.1) !important;
    margin-bottom: 30px !important;
}

/* Header Styling */
.header-container {
    display: flex;
    flex-direction: column;
    align-items: center;
    margin-bottom: 2rem;
    padding-top: 1rem;
}
.logo-svg {
    animation: float 3s ease-in-out infinite;
    margin-bottom: 10px;
}
h1.title-text {
    font-size: 32px !important;
    font-weight: 800 !important;
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    margin: 0 !important;
    text-align: center;
}
.tagline {
    text-align: center;
    color: rgba(255, 255, 255, 0.6);
    font-size: 16px;
    margin-top: 5px;
}

/* Upload Area Beautification */
div.svelte-116pqw2 {
    background: rgba(255,255,255,0.03) !important;
    border: 2px dashed rgba(102, 126, 234, 0.4) !important;
    border-radius: 16px !important;
    transition: all 0.3s ease;
}
div.svelte-116pqw2:hover {
    border-color: #667eea !important;
    background: rgba(102, 126, 234, 0.05) !important;
}

/* Button Exact Replica */
button.primary-btn {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
    color: white !important;
    border: none !important;
    padding: 18px 48px !important;
    border-radius: 16px !important;
    font-size: 18px !important;
    font-weight: 700 !important;
    cursor: pointer !important;
    box-shadow: 0 4px 20px rgba(102, 126, 234, 0.5) !important;
    width: 100% !important;
    transition: all 0.3s ease !important;
    margin-top: 20px !important;
}
button.primary-btn:hover {
    transform: translateY(-2px);
    box-shadow: 0 8px 30px rgba(102, 126, 234, 0.7) !important;
}

/* Result Cards Styling */
.result-header {
    display: flex;
    align-items: center;
    gap: 10px;
    margin-bottom: 15px;
    padding-bottom: 10px;
    border-bottom: 1px solid rgba(255,255,255,0.1);
}
.track-icon {
    width: 40px;
    height: 40px;
    border-radius: 10px;
    display: flex;
    align-items: center;
    justify-content: center;
    font-size: 20px;
}
.karaoke-icon { background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); }
.vocals-icon { background: linear-gradient(135deg, #f093fb 0%, #f5576c 100%); }

/* Animations */
@keyframes float {
    0%, 100% { transform: translateY(0); }
    50% { transform: translateY(-10px); }
}
@keyframes pulse {
    0%, 100% { opacity: 1; transform: scale(1); }
    50% { opacity: 0.8; transform: scale(0.95); }
}

footer { display: none !important; }
"""

def separate_audio(audio_file, progress=gr.Progress()):
    if audio_file is None:
        return None, None

    # Step 1: Initialization
    progress(0, desc="Initializing...")
    time.sleep(1)
    
    # Create unique job directory
    job_id = str(uuid.uuid4())
    output_dir = os.path.join("output", job_id)
    os.makedirs(output_dir, exist_ok=True)
    
    try:
        # Step 2: Running Demucs
        progress(0.2, desc="Separating Vocals (AI Processing)...")
        
        cmd = [
            "demucs",
            "--out", output_dir,
            "--two-stems", "vocals",
            "-n", "htdemucs",
            "--device", "cpu",
            audio_file
        ]
        
        # We can't track real-time Demucs progress easily in subprocess without parsing output
        # So we just run it. The user will see the 0.2 bar while waiting.
        result = subprocess.run(cmd, capture_output=True, text=True)
        
        if result.returncode != 0:
            raise Exception(f"Demucs failed: {result.stderr}")

        # Step 3: Finding files
        progress(0.7, desc="Processing Output Files...")
        
        model_output_dir = os.path.join(output_dir, "htdemucs")
        song_folders = os.listdir(model_output_dir)
        if not song_folders:
            raise Exception("No output folder found")
            
        song_folder = os.path.join(model_output_dir, song_folders[0])
        
        vocals_wav = os.path.join(song_folder, "vocals.wav")
        karaoke_wav = os.path.join(song_folder, "no_vocals.wav")
        
        # Step 4: Converting to MP3
        progress(0.8, desc="Converting to MP3 (Optimizing size)...")
        
        vocals_mp3 = vocals_wav.replace(".wav", ".mp3")
        karaoke_mp3 = karaoke_wav.replace(".wav", ".mp3")
        
        AudioSegment.from_wav(vocals_wav).export(vocals_mp3, format="mp3", bitrate="320k")
        AudioSegment.from_wav(karaoke_wav).export(karaoke_mp3, format="mp3", bitrate="320k")
        
        progress(1.0, desc="Done!")
        return karaoke_mp3, vocals_mp3

    except Exception as e:
        print(f"Error: {str(e)}")
        raise gr.Error(f"Processing failed: {str(e)}")

# --- HTML COMPONENTS ---
header_html = """
<div class="header-container">
    <svg class="logo-svg" width="48" height="48" viewBox="0 0 32 32" fill="none">
        <circle cx="16" cy="16" r="14" stroke="url(#gradient)" stroke-width="2"/>
        <path d="M12 10L12 22M16 8L16 24M20 12L20 20" stroke="url(#gradient)" stroke-width="2" stroke-linecap="round"/>
        <defs>
            <linearGradient id="gradient" x1="0" y1="0" x2="32" y2="32">
                <stop offset="0%" stop-color="#667eea" />
                <stop offset="100%" stop-color="#764ba2" />
            </linearGradient>
        </defs>
    </svg>
    <h1 class="title-text">Karaoke Extractor</h1>
    <p class="tagline">Create studio-quality instrumentals & acapellas in seconds</p>
</div>
"""

upload_icon_html = """
<div style="display: flex; flex-direction: column; align-items: center; margin-bottom: 20px;">
    <div style="margin-bottom: 15px;">
        <svg width="80" height="80" viewBox="0 0 64 64" fill="none" style="animation: pulse 2s ease-in-out infinite;">
            <circle cx="32" cy="32" r="30" stroke="url(#uploadGradient)" stroke-width="2" stroke-dasharray="4 4" opacity="0.4"/>
            <path d="M32 20V44M20 32L32 20L44 32" stroke="url(#uploadGradient)" stroke-width="3" stroke-linecap="round" stroke-linejoin="round"/>
            <defs>
                <linearGradient id="uploadGradient" x1="0" y1="0" x2="64" y2="64">
                    <stop offset="0%" stop-color="#667eea" />
                    <stop offset="100%" stop-color="#764ba2" />
                </linearGradient>
            </defs>
        </svg>
    </div>
    <h2 style="font-size: 24px; font-weight: 700; margin-bottom: 8px;">Upload Your Song</h2>
    <p style="color: rgba(255,255,255,0.6); font-size: 15px;">Support for MP3, WAV, FLAC, M4A</p>
</div>
"""

karaoke_header = """
<div class="result-header">
    <div class="track-icon karaoke-icon">🎸</div>
    <div>
        <h3 style="margin:0; font-size:18px;">Karaoke Track</h3>
        <p style="margin:0; font-size:12px; opacity:0.7;">Instrumental (No Vocals)</p>
    </div>
</div>
"""

vocals_header = """
<div class="result-header">
    <div class="track-icon vocals-icon">🎤</div>
    <div>
        <h3 style="margin:0; font-size:18px;">Vocals Track</h3>
        <p style="margin:0; font-size:12px; opacity:0.7;">Acapella (Vocals Only)</p>
    </div>
</div>
"""

# Build the UI
with gr.Blocks(css=custom_css, title="Karaoke Extractor") as app:
    
    gr.HTML(header_html)
    
    with gr.Group(elem_classes="glass-card"):
        gr.HTML(upload_icon_html)
        audio_input = gr.Audio(type="filepath", label="", sources=["upload"], container=False, interactive=True)
        process_btn = gr.Button("✨ Extract Tracks Now", elem_classes="primary-btn")
    
    # Results Area
    with gr.Group(elem_classes="glass-card", visible=True):
        gr.Markdown("<h2 style='text-align: center; margin-bottom: 20px;'>✨ Your Extracted Tracks</h2>")
        
        with gr.Row():
            with gr.Column():
                gr.HTML(karaoke_header)
                karaoke_output = gr.Audio(label="", type="filepath", interactive=False, show_label=False)
            
            with gr.Column():
                gr.HTML(vocals_header)
                vocals_output = gr.Audio(label="", type="filepath", interactive=False, show_label=False)

    process_btn.click(
        fn=separate_audio,
        inputs=[audio_input],
        outputs=[karaoke_output, vocals_output]
    )

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