import gradio as gr
import requests
import os
import json
import uuid
from datetime import datetime

# Load Suno API key from environment variable
SUNO_KEY = os.environ.get("SUNO_KEY", os.environ.get("SUNOKEY", ""))

# In production, use your actual public URL
# For Spaces, get from environment or use ngrok for local testing
SPACE_URL = os.environ.get("SPACE_URL", "https://your-username.hf.space")

# Store tasks with timestamps
tasks_db = {}

def generate_lyrics(prompt, callBackUrl=""):
    """Submit lyrics generation task with callback URL"""
    if not SUNO_KEY:
        return "❌ Error: SUNO_KEY environment variable not set"
    
    # Generate unique task ID
    task_id = str(uuid.uuid4())[:8]
    
    # Use provided callback URL or default to your Space URL
    callback_url = callBackUrl or f"{SPACE_URL}/callback"
    
    url = "https://api.sunoapi.org/api/v1/lyrics"
    headers = {
        "Authorization": f"Bearer {SUNO_KEY}",
        "Content-Type": "application/json"
    }
    
    payload = {
        "prompt": prompt,
        "callBackUrl": callback_url,
        "customTaskId": task_id  # Optional: track your own ID
    }
    
    try:
        resp = requests.post(url, headers=headers, json=payload, timeout=30)
        data = resp.json()
        
        if resp.status_code == 200 and data.get("code") == 200:
            api_task_id = data["data"]["taskId"]
            
            # Store task info
            tasks_db[task_id] = {
                "api_task_id": api_task_id,
                "prompt": prompt,
                "status": "submitted",
                "submitted_at": datetime.now().isoformat(),
                "callback_received": False,
                "lyrics": None
            }
            
            return f"""✅ Task Submitted!
Your Task ID: {task_id}
API Task ID: {api_task_id}
Callback URL: {callback_url}

📝 Prompt: {prompt}

⏳ Processing... The results will be sent to the callback URL.
You can also poll manually using your Task ID above."""
        
        else:
            error_msg = data.get("msg", "Unknown error")
            return f"❌ API Error: {error_msg} (Code: {data.get('code')})"
            
    except Exception as e:
        return f"❌ Error: {str(e)}"

def poll_task(task_id):
    """Manual polling as fallback"""
    if not SUNO_KEY:
        return "❌ Error: SUNO_KEY not configured"
    
    if task_id not in tasks_db:
        return "❌ Task ID not found. Please submit a task first."
    
    task_info = tasks_db[task_id]
    api_task_id = task_info["api_task_id"]
    
    url = f"https://api.sunoapi.org/api/v1/lyrics/details?taskId={api_task_id}"
    headers = {"Authorization": f"Bearer {SUNO_KEY}"}
    
    try:
        resp = requests.get(url, headers=headers, timeout=30)
        data = resp.json()
        
        if resp.status_code == 200 and data.get("code") == 200:
            task_data = data["data"]
            status = task_data.get("status", "unknown")
            tasks_db[task_id]["status"] = status
            
            if status == "completed" and "data" in task_data:
                lyrics_data = task_data["data"]
                tasks_db[task_id]["lyrics"] = lyrics_data
                tasks_db[task_id]["callback_received"] = True
                
                return format_lyrics(lyrics_data, task_id)
            elif status == "failed":
                return f"❌ Task failed: {task_data.get('error', 'Unknown error')}"
            else:
                return f"⏳ Status: {status}\nLast checked: {datetime.now().strftime('%H:%M:%S')}"
        else:
            return f"❌ Polling error: {data.get('msg', 'Unknown error')}"
            
    except Exception as e:
        return f"❌ Error: {str(e)}"

def format_lyrics(lyrics_data, task_id):
    """Format lyrics for display"""
    output = [f"🎵 **Task {task_id} - Generated Lyrics**", ""]
    
    for i, item in enumerate(lyrics_data, 1):
        title = item.get('title', f'Variant {i}')
        text = item.get('text', 'No lyrics generated')
        
        output.append(f"**Variant {i}: {title}**")
        output.append("```")
        output.append(text)
        output.append("```")
        output.append("---")
    
    output.append(f"\n✅ Generated at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    return "\n".join(output)

def list_tasks():
    """Show all submitted tasks"""
    if not tasks_db:
        return "No tasks submitted yet."
    
    output = ["📋 **Submitted Tasks:**", ""]
    for task_id, info in tasks_db.items():
        status_icon = "✅" if info["callback_received"] else "⏳"
        output.append(f"{status_icon} **{task_id}** - {info['status']}")
        output.append(f"   Prompt: {info['prompt'][:50]}...")
        output.append(f"   Submitted: {info['submitted_at']}")
        output.append("")
    
    return "\n".join(output)

# WEBHOOK ENDPOINT (for receiving callbacks)
def webhook_callback(request: gr.Request):
    """Handle incoming webhook from Suno API"""
    try:
        # Try to get JSON data
        data = request.json()
        
        if not data:
            # Try form data
            data = dict(request.form)
        
        print(f"📥 Received webhook: {json.dumps(data, indent=2)}")
        
        # Extract task info from webhook
        # Suno API format might vary - adjust based on actual response
        if "data" in data and "taskId" in data["data"]:
            api_task_id = data["data"]["taskId"]
            status = data["data"].get("status", "unknown")
            
            # Find our task by API task ID
            for task_id, task_info in tasks_db.items():
                if task_info["api_task_id"] == api_task_id:
                    task_info["status"] = status
                    task_info["callback_received"] = True
                    
                    if status == "completed" and "data" in data["data"]:
                        task_info["lyrics"] = data["data"]["data"]
                        print(f"✅ Lyrics received for task {task_id}")
                    
                    return {"status": "success", "message": f"Updated task {task_id}"}
        
        return {"status": "error", "message": "Task not found"}
        
    except Exception as e:
        print(f"❌ Webhook error: {e}")
        return {"status": "error", "message": str(e)}

# Gradio Interface
with gr.Blocks(theme=gr.themes.Soft(), title="Suno Lyrics Generator") as app:
    gr.Markdown("# 🎵 Suno AI Lyrics Generator")
    gr.Markdown("Generate song lyrics with webhook support")
    
    with gr.Tabs():
        with gr.TabItem("🎤 Generate"):
            with gr.Row():
                with gr.Column():
                    prompt = gr.Textbox(
                        label="Lyrics Prompt",
                        placeholder="A romantic ballad about stargazing...",
                        lines=3
                    )
                    
                    # Optional: Custom callback URL
                    callback_url = gr.Textbox(
                        label="Callback URL (Optional)",
                        value=f"{SPACE_URL}/callback",
                        info="Where Suno should send results. Leave as default for Spaces."
                    )
                    
                    submit_btn = gr.Button("✨ Generate Lyrics", variant="primary")
                    
                    gr.Markdown("### ℹ️ Instructions:")
                    gr.Markdown("""
                    1. Enter your lyrics prompt
                    2. Click Generate
                    3. Save your Task ID
                    4. Check status in the "Poll Tasks" tab
                    5. Results will arrive via webhook automatically
                    """)
                
                with gr.Column():
                    output = gr.Textbox(
                        label="Submission Result",
                        lines=10,
                        interactive=False
                    )
            
            submit_btn.click(
                generate_lyrics,
                inputs=[prompt, callback_url],
                outputs=output
            )
        
        with gr.TabItem("🔄 Poll Tasks"):
            with gr.Row():
                with gr.Column():
                    task_id_input = gr.Textbox(
                        label="Your Task ID",
                        placeholder="Enter the Task ID from generation step"
                    )
                    poll_btn = gr.Button("🔍 Check Status", variant="primary")
                    
                    gr.Markdown("---")
                    refresh_btn = gr.Button("📋 List All Tasks")
                    tasks_list = gr.Textbox(label="All Tasks", lines=10)
                
                with gr.Column():
                    poll_result = gr.Textbox(
                        label="Task Status",
                        lines=15,
                        interactive=False
                    )
            
            poll_btn.click(
                poll_task,
                inputs=task_id_input,
                outputs=poll_result
            )
            
            refresh_btn.click(
                list_tasks,
                inputs=None,
                outputs=tasks_list
            )
        
        with gr.TabItem("⚙️ Webhook Info"):
            gr.Markdown("### 🌐 Webhook Configuration")
            gr.Markdown(f"""
            **Your Webhook URL:** `{SPACE_URL}/callback`
            
            **For Local Development:**
            1. Use [ngrok](https://ngrok.com/): `ngrok http 7860`
            2. Update callback URL: `https://your-ngrok-url.ngrok.io/callback`
            
            **For Hugging Face Spaces:**
            - The URL above should work automatically
            - Make sure your Space is public or has network access
            """)
            
            webhook_status = gr.Textbox(
                label="Last Webhook Status",
                value="No webhooks received yet",
                lines=5
            )
    
    # Register webhook endpoint
    # Note: In production, you'd set up proper route handling
    # For Gradio, we can simulate with a POST endpoint
    app.post("/callback")(webhook_callback)

# Launch configuration
if __name__ == "__main__":
    # For local testing with webhooks
    import socket
    
    # Get local IP for ngrok compatibility
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        s.connect(("8.8.8.8", 80))
        local_ip = s.getsockname()[0]
        s.close()
        print(f"🌐 Local IP: {local_ip}")
        print(f"🌐 Webhook URL: http://{local_ip}:7860/callback")
    except:
        pass
    
    app.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True
    )