import gradio as gr
import requests
import os
import time
import json
from datetime import datetime
import re

# Suno API key
SUNO_KEY = os.environ.get("SunoKey", "")
if not SUNO_KEY:
    print("⚠️ SunoKey not set in environment variables!")

# API endpoints
API_BASE = "https://api.sunoapi.org"
CREDIT_URL = f"{API_BASE}/api/v1/generate/credit"
GENERATION_RECORD_URL = f"{API_BASE}/api/v1/generate/record-info"
VOCAL_SEPARATION_URL = f"{API_BASE}/api/v1/vocal-removal/generate"
SEPARATION_RECORD_URL = f"{API_BASE}/api/v1/vocal-removal/record-info"

def check_credits():
    """Check available credits"""
    try:
        response = requests.get(CREDIT_URL, headers={"Authorization": f"Bearer {SUNO_KEY}"}, timeout=10)
        if response.status_code == 200:
            data = response.json()
            if data.get("code") == 200:
                credits = data.get("data", {}).get("credit_balance", 0)
                return f"✅ Available credits: **{credits}**"
        return "⚠️ Could not retrieve credits"
    except:
        return "⚠️ Credit check failed"

def get_music_tasks():
    """Get list of recent music generation tasks"""
    if not SUNO_KEY:
        return "❌ API key not configured", []
    
    try:
        response = requests.get(
            GENERATION_RECORD_URL,
            headers={"Authorization": f"Bearer {SUNO_KEY}"},
            params={"page": 1, "pageSize": 20},
            timeout=30
        )
        
        if response.status_code == 200:
            data = response.json()
            if data.get("code") == 200:
                tasks = data.get("data", {}).get("data", [])
                
                if not tasks:
                    return "📭 No music generation tasks found. Please generate music first.", []
                
                task_list = []
                for task in tasks:
                    task_id = task.get("taskId", "")
                    status = task.get("status", "")
                    audio_urls = task.get("audioUrl", [])
                    
                    # Get audio IDs if available
                    audio_info = []
                    if isinstance(audio_urls, list):
                        for i, url in enumerate(audio_urls[:2]):  # First 2 variants
                            # Extract audio ID from URL or use placeholder
                            audio_id = extract_audio_id(url) or f"variant-{i+1}"
                            audio_info.append(f"Audio {i+1}: {audio_id}")
                    
                    task_info = f"Task: `{task_id[:8]}...` | Status: {status}"
                    if audio_info:
                        task_info += f" | {', '.join(audio_info)}"
                    
                    task_list.append((task_id, task_info))
                
                return f"✅ Found {len(task_list)} recent tasks", task_list
            else:
                return f"❌ API error: {data.get('msg')}", []
        else:
            return f"❌ HTTP error: {response.status_code}", []
            
    except Exception as e:
        return f"❌ Error: {str(e)}", []

def extract_audio_id(url):
    """Extract audio ID from URL"""
    if not url:
        return None
    
    # Try to find UUID pattern in URL
    uuid_pattern = r'[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}'
    match = re.search(uuid_pattern, url, re.IGNORECASE)
    if match:
        return match.group(0)
    
    # Try to find in common URL patterns
    patterns = [
        r'audioId=([^&]+)',
        r'audio/([^/]+)',
        r'id=([^&]+)'
    ]
    
    for pattern in patterns:
        match = re.search(pattern, url)
        if match:
            return match.group(1)
    
    return None

def get_audio_ids_from_task(task_id):
    """Get audio IDs for a specific task"""
    if not task_id:
        return "❌ Please select a task first", []
    
    try:
        response = requests.get(
            GENERATION_RECORD_URL,
            headers={"Authorization": f"Bearer {SUNO_KEY}"},
            params={"taskId": task_id},
            timeout=30
        )
        
        if response.status_code == 200:
            data = response.json()
            if data.get("code") == 200:
                task_data = data.get("data", {})
                audio_urls = task_data.get("audioUrl", [])
                
                if isinstance(audio_urls, list) and audio_urls:
                    audio_options = []
                    for i, url in enumerate(audio_urls):
                        audio_id = extract_audio_id(url) or f"audio-{i+1}"
                        display_name = f"Audio {i+1}"
                        if audio_id:
                            display_name += f" (ID: {audio_id[:8]}...)"
                        audio_options.append((url, display_name))
                    
                    return f"✅ Found {len(audio_options)} audio tracks", audio_options
                else:
                    return "⚠️ No audio tracks found for this task", []
            else:
                return f"❌ API error: {data.get('msg')}", []
        else:
            return f"❌ HTTP error: {response.status_code}", []
            
    except Exception as e:
        return f"❌ Error: {str(e)}", []

def separate_vocals(task_id, audio_url, separation_type):
    """Separate vocals and instruments from selected audio"""
    if not SUNO_KEY:
        yield "❌ Error: SunoKey not configured"
        return
    
    if not task_id or not audio_url:
        yield "❌ Error: Please select both Task and Audio"
        return
    
    # Extract audio ID from URL
    audio_id = extract_audio_id(audio_url)
    if not audio_id:
        yield "❌ Could not extract Audio ID from the selected audio. Please check the URL format."
        return
    
    # Validate separation type
    if separation_type not in ["separate_vocal", "split_stem"]:
        yield "❌ Error: Invalid separation type"
        return
    
    # Submit separation task
    try:
        resp = requests.post(
            VOCAL_SEPARATION_URL,
            json={
                "taskId": task_id,
                "audioId": audio_id,
                "type": separation_type,
                "callBackUrl": "http://dummy.com/callback"
            },
            headers={
                "Authorization": f"Bearer {SUNO_KEY}",
                "Content-Type": "application/json"
            },
            timeout=30
        )
        
        if resp.status_code != 200:
            yield f"❌ Submission failed: HTTP {resp.status_code}"
            return
        
        data = resp.json()
        if data.get("code") != 200:
            yield f"❌ API error: {data.get('msg', 'Unknown')}"
            return
        
        separation_task_id = data["data"]["taskId"]
        yield f"✅ **Separation Task Submitted!**\n\n"
        yield f"**Separation Task ID:** `{separation_task_id}`\n"
        yield f"**Original Task ID:** `{task_id}`\n"
        yield f"**Audio ID:** `{audio_id}`\n"
        yield f"**Separation Type:** {separation_type}\n\n"
        yield f"⏳ **Processing...** (Usually takes 30-120 seconds)\n\n"
        
        # Poll for results
        for attempt in range(40):  # 40 attempts * 5 seconds = 200 seconds max
            time.sleep(5)
            
            try:
                check = requests.get(
                    SEPARATION_RECORD_URL,
                    headers={"Authorization": f"Bearer {SUNO_KEY}"},
                    params={"taskId": separation_task_id},
                    timeout=30
                )
                
                if check.status_code == 200:
                    check_data = check.json()
                    status = check_data["data"].get("successFlag", "PENDING")
                    
                    if status == "SUCCESS":
                        # Success! Extract separation results
                        separation_info = check_data["data"].get("response", {})
                        
                        # Format output based on separation type
                        output = "🎵 **Separation Complete!**\n\n"
                        
                        if separation_type == "separate_vocal":
                            output += "## 🎤 2-Stem Separation\n\n"
                            output += f"**🎤 Vocals:** [Download]({separation_info.get('vocalUrl', 'N/A')})\n"
                            output += f"**🎵 Instrumental:** [Download]({separation_info.get('instrumentalUrl', 'N/A')})\n"
                            if separation_info.get('originUrl'):
                                output += f"**📁 Original:** [Download]({separation_info.get('originUrl')})\n"
                            
                        elif separation_type == "split_stem":
                            output += "## 🎛️ 12-Stem Separation\n\n"
                            
                            stems = [
                                ("🎤 Vocals", separation_info.get('vocalUrl')),
                                ("🎤 Backing Vocals", separation_info.get('backingVocalsUrl')),
                                ("🥁 Drums", separation_info.get('drumsUrl')),
                                ("🎸 Bass", separation_info.get('bassUrl')),
                                ("🎸 Guitar", separation_info.get('guitarUrl')),
                                ("🎹 Keyboard", separation_info.get('keyboardUrl')),
                                ("🎻 Strings", separation_info.get('stringsUrl')),
                                ("🎺 Brass", separation_info.get('brassUrl')),
                                ("🎷 Woodwinds", separation_info.get('woodwindsUrl')),
                                ("🪇 Percussion", separation_info.get('percussionUrl')),
                                ("🎹 Synth", separation_info.get('synthUrl')),
                                ("🎛️ FX/Other", separation_info.get('fxUrl')),
                                ("📁 Original", separation_info.get('originUrl'))
                            ]
                            
                            for stem_name, stem_url in stems:
                                if stem_url:
                                    output += f"**{stem_name}:** [Download]({stem_url})\n"
                        
                        output += f"\n⏱️ **Processing Time:** {(attempt + 1) * 5} seconds\n"
                        output += f"📊 **Separation Task ID:** `{separation_task_id}`\n"
                        output += f"⚠️ **Note:** Download links expire in 14 days\n"
                        
                        yield output
                        return
                    
                    elif status in ["CREATE_TASK_FAILED", "GENERATE_AUDIO_FAILED", "CALLBACK_EXCEPTION"]:
                        error = check_data["data"].get("errorMessage", "Unknown error")
                        yield f"❌ **Separation failed:** {status}\n\n"
                        yield f"**Error:** {error}\n"
                        return
                    
                    else:
                        # PENDING or other statuses
                        yield (f"⏳ **Status:** {status}\n"
                               f"📊 **Progress:** {attempt + 1}/40 attempts\n"
                               f"⏱️ **Elapsed:** {(attempt + 1) * 5} seconds\n\n"
                               f"Still processing...")
                else:
                    yield f"⚠️ **Check error:** HTTP {check.status_code}\n"
            
            except Exception as e:
                yield f"⚠️ **Error checking status:** {str(e)}\n"
        
        yield "⏰ **Timeout after 200 seconds.** Try checking the task manually later.\n"
        yield f"**Separation Task ID:** `{separation_task_id}`\n"
    
    except Exception as e:
        yield f"❌ **Error submitting task:** {str(e)}"

def refresh_tasks():
    """Refresh the list of available tasks"""
    status, tasks = get_music_tasks()
    task_options = [("", "Select a task...")] + tasks if tasks else [("", "No tasks found")]
    return status, gr.Dropdown(choices=[opt[1] for opt in task_options], value=task_options[0][1]), task_options

# Create the app
with gr.Blocks(title="Suno Stem Separator", theme="soft") as app:
    gr.Markdown("# 🎵 Suno AI Stem Separator")
    gr.Markdown("Separate vocals and instruments from your Suno AI generated music")
    
    # Credit display at top
    credit_display = gr.Markdown(check_credits(), elem_id="credits")
    
    with gr.Tabs():
        with gr.TabItem("🎵 Step 1: Select Your Music"):
            gr.Markdown("""
            ### Find your generated music tracks
            
            1. Click **Refresh Tasks** to load your recent Suno AI music generations
            2. Select a task from the dropdown
            3. The system will automatically find available audio tracks
            """)
            
            with gr.Row():
                refresh_btn = gr.Button("🔄 Refresh Tasks", variant="secondary")
                task_status = gr.Markdown("Click 'Refresh Tasks' to load your music tasks")
            
            with gr.Row():
                with gr.Column(scale=1):
                    task_dropdown = gr.Dropdown(
                        label="Select Music Generation Task",
                        choices=["Select a task..."],
                        value="Select a task...",
                        interactive=True
                    )
                    
                    gr.Markdown("""
                    **What is a Task ID?**
                    - When you generate music with Suno AI, each generation gets a unique Task ID
                    - This ID is required to separate vocals from your music
                    - Select your task from the dropdown above
                    """)
                
                with gr.Column(scale=1):
                    audio_status = gr.Markdown("Select a task to see available audio tracks")
                    audio_dropdown = gr.Dropdown(
                        label="Select Audio Track to Separate",
                        choices=["First select a task..."],
                        value="First select a task...",
                        interactive=True
                    )
                    
                    gr.Markdown("""
                    **What is an Audio ID?**
                    - Each music generation can have multiple audio variants
                    - The Audio ID identifies which specific track to process
                    - Usually you'll have 2-4 audio tracks per generation
                    """)
            
            # Hidden storage for task data
            task_store = gr.State([])  # Stores (task_id, display_name) pairs
        
        with gr.TabItem("🎛️ Step 2: Configure Separation"):
            gr.Markdown("""
            ### Configure how you want to separate the audio
            
            **Choose separation type:**
            - **🎤 Vocal Separation (1 credit)**: Separate vocals from instrumental
            - **🎛️ Full Stem Separation (5 credits)**: Split into 12 individual instrument stems
            
            **Stem types in full separation:**
            - Vocals, Backing Vocals, Drums, Bass, Guitar
            - Keyboard, Strings, Brass, Woodwinds, Percussion
            - Synthesizer, FX/Other
            """)
            
            separation_type = gr.Radio(
                label="Separation Type",
                choices=[
                    ("🎤 Vocal Separation (2 stems, 1 credit)", "separate_vocal"),
                    ("🎛️ Full Stem Separation (12 stems, 5 credits)", "split_stem")
                ],
                value="separate_vocal",
                info="Choose how detailed the separation should be"
            )
            
            with gr.Row():
                selected_task_display = gr.Textbox(
                    label="Selected Task ID",
                    interactive=False,
                    value="No task selected"
                )
                selected_audio_display = gr.Textbox(
                    label="Selected Audio ID",
                    interactive=False,
                    value="No audio selected"
                )
            
            separate_btn = gr.Button("🚀 Start Separation", variant="primary", scale=1)
            
            gr.Markdown("""
            **⚠️ Important Notes:**
            - Each separation consumes credits (shown at top of page)
            - Processing takes 30-120 seconds
            - Download links expire after 14 days
            - You can reuse the Separation Task ID to download files later
            """)
        
        with gr.TabItem("📥 Step 3: Get Results"):
            output = gr.Markdown(
                label="Separation Results",
                value="Your separated stems will appear here once processing is complete..."
            )
            
            gr.Markdown("""
            **What to expect:**
            1. After clicking **Start Separation**, you'll get a Separation Task ID
            2. The system will automatically check progress every 5 seconds
            3. When complete, you'll see download links for all stems
            4. Links remain active for **14 days** - download promptly!
            
            **Need to check a previous separation?**
            Use the **Check Existing Task** tab (coming soon!)
            """)
    
    gr.Markdown("---")
    gr.Markdown(
        """
        <div style="text-align: center; padding: 20px;">
            <p>Powered by <a href="https://suno.ai" target="_blank">Suno AI</a> • 
            <a href="https://sunoapi.org" target="_blank">Suno API</a> • 
            <a href="https://docs.sunoapi.org" target="_blank">API Docs</a></p>
            <p><small>This tool helps you separate vocals and instruments from Suno AI generated music.</small></p>
        </div>
        """,
        elem_id="footer"
    )
    
    # Event handlers
    def on_task_select(selected_display, task_store_data):
        """When a task is selected from dropdown"""
        # Find the actual task_id from the display name
        for task_id, display_name in task_store_data:
            if display_name == selected_display:
                # Get audio IDs for this task
                status, audio_options = get_audio_ids_from_task(task_id)
                
                # Update audio dropdown
                audio_choices = [("", "Select audio...")] + audio_options if audio_options else [("", "No audio found")]
                
                return (
                    status,
                    gr.Dropdown(choices=[opt[1] for opt in audio_choices], value=audio_choices[0][1]),
                    task_id[:12] + "..." if len(task_id) > 12 else task_id,
                    "Select audio above..."
                )
        
        return "❌ Task not found", gr.Dropdown(choices=["Task not found"], value="Task not found"), "Error", "Error"
    
    def on_audio_select(selected_audio_display, audio_options_data):
        """When an audio is selected"""
        # Extract audio ID from the selected option
        # The display format is "Audio 1 (ID: abc123...)" or similar
        import re
        match = re.search(r'ID:\s*([^\s)]+)', selected_audio_display)
        audio_id = match.group(1) if match else selected_audio_display
        
        return audio_id
    
    # Connect events
    refresh_btn.click(
        refresh_tasks,
        outputs=[task_status, task_dropdown, task_store]
    )
    
    task_dropdown.change(
        on_task_select,
        inputs=[task_dropdown, task_store],
        outputs=[audio_status, audio_dropdown, selected_task_display, selected_audio_display]
    )
    
    audio_dropdown.change(
        lambda x: x,  # Simple pass-through for now
        inputs=[audio_dropdown],
        outputs=[selected_audio_display]
    )
    
    separate_btn.click(
        separate_vocals,
        inputs=[selected_task_display, audio_dropdown, separation_type],
        outputs=[output]
    )

if __name__ == "__main__":
    print("🚀 Starting Suno Stem Separator")
    print(f"🔑 SunoKey: {'✅ Set' if SUNO_KEY else '❌ Not set'}")
    print("🌐 Open your browser to: http://localhost:7860")
    print("\n💡 Instructions:")
    print("1. Make sure you have generated music with Suno AI")
    print("2. Refresh tasks to load your music generations")
    print("3. Select a task and audio track")
    print("4. Choose separation type and start processing")
    
    app.launch(server_name="0.0.0.0", server_port=7860, share=False)