import gradio as gr
import requests
import os
import time
import json

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

def get_audio_info(task_id):
    """Get audio information from Suno task ID"""
    if not SUNO_KEY:
        return "❌ Error: SunoKey not configured in environment variables", []
    
    if not task_id.strip():
        return "❌ Error: Please enter Task ID", []
    
    try:
        resp = requests.get(
            "https://api.sunoapi.org/api/v1/generate/record-info",
            headers={"Authorization": f"Bearer {SUNO_KEY}"},
            params={"taskId": task_id.strip()},
            timeout=30
        )
        
        if resp.status_code != 200:
            return f"❌ Request failed: HTTP {resp.status_code}", []
        
        data = resp.json()
        if data.get("code") != 200:
            return f"❌ API error: {data.get('msg', 'Unknown')}", []
        
        # Check if task is complete
        status = data["data"].get("status", "PENDING")
        if status != "SUCCESS":
            return f"⏳ Task status: {status}. Please wait for generation to complete.", []
        
        # Parse sunoData to get audio options
        suno_data = data["data"]["response"].get("sunoData", [])
        if not suno_data:
            return "❌ No audio data found in response", []
        
        # Create audio options for dropdown
        audio_options = []
        for i, audio in enumerate(suno_data):
            audio_id = audio.get("id", "")
            prompt = audio.get("prompt", "No prompt")
            title = audio.get("title", "Untitled")
            duration = audio.get("duration", 0)
            
            # Truncate long prompts for display
            display_prompt = prompt[:50] + "..." if len(prompt) > 50 else prompt
            
            # Create display text and store full data
            display_text = f"Track {i+1}: {title} ({duration:.1f}s) - {display_prompt}"
            audio_options.append((display_text, audio_id, audio))
        
        if not audio_options:
            return "❌ No valid audio tracks found", []
        
        # Format output message
        output = f"✅ **Task Found!**\n"
        output += f"**Task ID:** `{task_id}`\n"
        output += f"**Status:** {status}\n"
        output += f"**Found {len(audio_options)} audio track(s):**\n\n"
        
        for i, (display_text, audio_id, audio) in enumerate(audio_options):
            output += f"**Track {i+1}:**\n"
            output += f"- **ID:** `{audio_id}`\n"
            output += f"- **Title:** {audio.get('title', 'Untitled')}\n"
            output += f"- **Prompt:** {audio.get('prompt', 'No prompt')[:100]}...\n"
            output += f"- **Duration:** {audio.get('duration', 0):.1f}s\n"
            if audio.get('audioUrl'):
                output += f"- **Audio:** [Listen]({audio.get('audioUrl')})\n"
            output += "\n"
        
        output += "👇 **Select a track from the dropdown below to separate stems**"
        
        return output, audio_options
        
    except Exception as e:
        return f"❌ Error: {str(e)}", []

def separate_vocals(task_id, audio_id, separation_type, audio_data=None):
    """Separate vocals and instruments from Suno tracks"""
    if not SUNO_KEY:
        yield "❌ Error: SunoKey not configured in environment variables"
        return
    
    if not task_id.strip() or not audio_id.strip():
        yield "❌ Error: Please enter both Task ID and Audio ID"
        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(
            "https://api.sunoapi.org/api/v1/vocal-removal/generate",
            json={
                "taskId": task_id,
                "audioId": audio_id,
                "type": separation_type,
                "callBackUrl": "https://1hit.no/callback.php"  # Required but not used for polling
            },
            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"✅ **Submitted!**\nSeparation Task ID: `{separation_task_id}`\n\n⏳ Processing separation...\n"
        
        # Poll for results
        for attempt in range(40):  # 40 attempts * 5 seconds = 200 seconds max
            time.sleep(5)
            
            try:
                check = requests.get(
                    "https://api.sunoapi.org/api/v1/vocal-removal/record-info",
                    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("status", "PENDING")
                    
                    if status == "SUCCESS":
                        # Success! Extract separation results
                        separation_info = check_data["data"].get("vocal_removal_info", {})
                        
                        if not separation_info:
                            yield "✅ Completed but no separation results found"
                            return
                        
                        # 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('vocal_url', 'N/A')})\n"
                            output += f"**Instrumental:** [Download]({separation_info.get('instrumental_url', 'N/A')})\n"
                            if separation_info.get('origin_url'):
                                output += f"**Original:** [Download]({separation_info.get('origin_url')})\n"
                            
                        elif separation_type == "split_stem":
                            output += "## 12-Stem Separation\n\n"
                            
                            stems = [
                                ("Vocals", separation_info.get('vocal_url')),
                                ("Backing Vocals", separation_info.get('backing_vocals_url')),
                                ("Drums", separation_info.get('drums_url')),
                                ("Bass", separation_info.get('bass_url')),
                                ("Guitar", separation_info.get('guitar_url')),
                                ("Keyboard", separation_info.get('keyboard_url')),
                                ("Strings", separation_info.get('strings_url')),
                                ("Brass", separation_info.get('brass_url')),
                                ("Woodwinds", separation_info.get('woodwinds_url')),
                                ("Percussion", separation_info.get('percussion_url')),
                                ("Synth", separation_info.get('synth_url')),
                                ("FX/Other", separation_info.get('fx_url')),
                                ("Instrumental", separation_info.get('instrumental_url')),
                                ("Original", separation_info.get('origin_url'))
                            ]
                            
                            for stem_name, stem_url in stems:
                                if stem_url:
                                    output += f"**{stem_name}:** [Download]({stem_url})\n"
                        
                        output += f"\n⏱️ Processed in about {(attempt + 1) * 5} seconds\n"
                        output += f"⚠️ **Note:** Download links expire in 14 days"
                        
                        yield output
                        return
                    
                    elif status == "FAILED":
                        error = check_data["data"].get("errorMessage", "Unknown error")
                        yield f"❌ Separation failed: {error}"
                        return
                    
                    else:
                        # PENDING or PROCESSING
                        yield (f"⏳ Status: {status}\n"
                               f"Attempt: {attempt + 1}/40\n"
                               f"Separation Task ID: `{separation_task_id}`\n\n"
                               f"Processing... (Usually takes 30-120 seconds)")
                else:
                    yield f"⚠️ Check error: HTTP {check.status_code}"
            
            except Exception as e:
                yield f"⚠️ Error checking status: {str(e)}"
        
        yield "⏰ Timeout after 200 seconds. Try checking again later."
    
    except Exception as e:
        yield f"❌ Error: {str(e)}"

# Create the app
with gr.Blocks(title="Suno Stem Separator", theme="soft") as app:
    gr.Markdown("# 🎵 Suno Stem Separator")
    gr.Markdown("Separate Suno AI tracks into vocal and instrument stems")
    
    with gr.Row():
        with gr.Column(scale=1):
            # Step 1: Enter Task ID
            task_id_input = gr.Textbox(
                label="Task ID",
                placeholder="Example: 5c79****be8e",
                info="Enter your Suno generation task ID",
                elem_id="task_id_input"
            )
            
            get_audio_btn = gr.Button("🔍 Find Audio Tracks", variant="secondary")
            
            # Will be updated after getting audio info
            audio_dropdown = gr.Dropdown(
                label="Select Audio Track to Separate",
                choices=[],
                info="Choose which audio track to process",
                visible=False,
                interactive=True
            )
            
            # Hidden store for audio data
            audio_data_store = gr.State(value=None)
            
            separation_type = gr.Radio(
                label="Separation Type",
                choices=[
                    ("separate_vocal (2 stems - 1 credit)", "separate_vocal"),
                    ("split_stem (12 stems - 5 credits)", "split_stem")
                ],
                value="separate_vocal",
                info="Choose separation mode",
                visible=False
            )
            
            separate_btn = gr.Button("🎵 Separate Stems", variant="primary", visible=False)
            
            gr.Markdown("""**Workflow:**
            1. Enter Task ID → Click "Find Audio Tracks"
            2. Select which audio track to separate
            3. Choose separation type
            4. Click "Separate Stems"
            5. Wait 30-120 seconds
            6. Get download links for each stem
            
            **Separation types:**
            - 🎤 **separate_vocal**: Vocals + Instrumental (2 stems, 1 credit)
            - 🎛️ **split_stem**: 12 detailed stems (5 credits)
            
            ⚠️ **Important:** 
            - Each request consumes credits
            - Links expire in 14 days
            """)
        
        with gr.Column(scale=2):
            status_output = gr.Markdown(
                label="Status",
                value="Enter a Task ID and click 'Find Audio Tracks'..."
            )
            
            results_output = gr.Markdown(
                label="Separation Results",
                visible=False
            )
    
    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 Docs</a> • 
            <a href="https://docs.sunoapi.org/separate-vocals" target="_blank">Stem Separation Guide</a></p>
            <p><small>This app uses the Suno API to separate tracks into individual stems.</small></p>
        </div>
        """,
        elem_id="footer"
    )
    
    # Step 1: Get audio info when button is clicked
    def on_get_audio(task_id):
        if not task_id:
            return (
                "❌ Please enter a Task ID",
                gr.Dropdown(choices=[], visible=False),
                gr.Radio(visible=False),
                gr.Button(visible=False),
                gr.Markdown(visible=False),
                None
            )
        
        output, audio_options = get_audio_info(task_id)
        
        if not audio_options:
            # No audio found, show error
            return (
                output,
                gr.Dropdown(choices=[], visible=False),
                gr.Radio(visible=False),
                gr.Button(visible=False),
                gr.Markdown(visible=False),
                None
            )
        
        # Create choices for dropdown
        choices = [(display_text, audio_id) for display_text, audio_id, _ in audio_options]
        
        # Store full audio data for later use
        audio_data_map = {audio_id: audio for _, audio_id, audio in audio_options}
        
        return (
            output,
            gr.Dropdown(choices=choices, value=choices[0][1] if choices else None, visible=True),
            gr.Radio(visible=True),
            gr.Button(visible=True),
            gr.Markdown(visible=False),
            audio_data_map
        )
    
    # Step 2: Separate stems when button is clicked
    def on_separate(task_id, audio_id, separation_type, audio_data_store):
        # Use the stored task_id from the first step
        yield from separate_vocals(task_id, audio_id, separation_type)
    
    # Connect events
    get_audio_btn.click(
        fn=on_get_audio,
        inputs=[task_id_input],
        outputs=[status_output, audio_dropdown, separation_type, separate_btn, results_output, audio_data_store]
    )
    
    separate_btn.click(
        fn=on_separate,
        inputs=[task_id_input, audio_dropdown, separation_type, audio_data_store],
        outputs=[results_output]
    ).then(
        lambda: gr.Markdown(visible=True),
        outputs=[results_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")
    app.launch(server_name="0.0.0.0", server_port=7860, share=False)