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 generate_lyrics(prompt):
    """Final working lyrics generator"""
    if not SUNO_KEY:
        yield "❌ Error: SunoKey not configured in environment variables"
        return
    
    if not prompt.strip():
        yield "❌ Error: Please enter a prompt"
        return
    
    # Submit task
    try:
        resp = requests.post(
            "https://api.sunoapi.org/api/v1/lyrics",
            json={
                "prompt": prompt,
                "callBackUrl": "http://dummy.com/callback"  # Required but not used
            },
            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
        
        task_id = data["data"]["taskId"]
        yield f"✅ **Submitted!**\nTask ID: `{task_id}`\n\n⏳ Waiting for lyrics...\n"
        
        # Poll for results
        for attempt in range(30):  # 30 attempts * 5 seconds = 150 seconds max
            time.sleep(5)
            
            try:
                check = requests.get(
                    "https://api.sunoapi.org/api/v1/lyrics/record-info",
                    headers={"Authorization": f"Bearer {SUNO_KEY}"},
                    params={"taskId": 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 lyrics
                        response_data = check_data["data"].get("response", {})
                        
                        if isinstance(response_data, str):
                            # Sometimes response is a JSON string
                            try:
                                response_data = json.loads(response_data.replace('null', 'None'))
                            except:
                                response_data = {"data": []}
                        
                        lyrics_list = response_data.get("data", [])
                        
                        if lyrics_list:
                            output = "🎵 **Lyrics Generated Successfully!**\n\n"
                            
                            for i, lyric in enumerate(lyrics_list, 1):
                                title = lyric.get('title', f'Variant {i}')
                                text = lyric.get('text', 'No lyrics')
                                
                                output += f"## Variant {i}: {title}\n"
                                output += "```\n"
                                output += text
                                output += "\n```\n"
                                output += "---\n\n"
                            
                            output += f"⏱️ Generated in about {(attempt + 1) * 5} seconds"
                            yield output
                        else:
                            yield "✅ Completed but no lyrics found in response"
                        return
                    
                    elif status == "FAILED":
                        error = check_data["data"].get("errorMessage", "Unknown error")
                        yield f"❌ Task failed: {error}"
                        return
                    
                    else:
                        # PENDING or PROCESSING
                        yield (f"⏳ Status: {status}\n"
                               f"Attempt: {attempt + 1}/30\n"
                               f"Task ID: `{task_id}`\n\n"
                               f"Still processing... (Usually takes 30-90 seconds)")
                else:
                    yield f"⚠️ Check error: HTTP {check.status_code}"
            
            except Exception as e:
                yield f"⚠️ Error checking status: {str(e)}"
        
        yield "⏰ Timeout after 150 seconds. Try checking again later."
    
    except Exception as e:
        yield f"❌ Error: {str(e)}"

# Create the app
with gr.Blocks(title="Suno Lyrics Generator", theme="soft") as app:
    gr.Markdown("# 🎵 Suno Lyrics Generator")
    gr.Markdown("Generate song lyrics using Suno AI")
    
    with gr.Row():
        with gr.Column(scale=1):
            prompt = gr.Textbox(
                label="Lyrics Prompt",
                placeholder="Example: A happy song about sunshine and rainbows",
                lines=3
            )
            btn = gr.Button("🎵 Generate Lyrics", variant="primary", scale=1)
            
            gr.Markdown("""
            **How it works:**
            1. Enter your lyrics idea
            2. Click Generate
            3. Wait 30-90 seconds
            4. Get 2 lyric variants
            
            **Status messages:**
            - ✅ Submitted = Task accepted
            - ⏳ PENDING/PROCESSING = Generating
            - 🎵 Success = Lyrics ready!
            """)
        
        with gr.Column(scale=2):
            output = gr.Markdown(
                label="Results",
                value="Your generated lyrics will appear here..."
            )
    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://github.com" target="_blank">GitHub</a></p>
            <p><small>This app uses the Suno API to generate AI-powered lyrics.</small></p>
        </div>
        """,
        elem_id="footer"
    )
    
    btn.click(generate_lyrics, prompt, output)

if __name__ == "__main__":
    print("🚀 Starting Suno Lyrics Generator")
    print(f"🔑 SunoKey: {'✅ Set' if SUNO_KEY else '❌ Not set'}")
    app.launch(server_name="0.0.0.0", server_port=7860, share=False)