import gradio as gr
import requests
import os
import json
import uuid
import time
import threading
from datetime import datetime
from typing import Dict, Optional

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

# Debug mode
DEBUG = True

# Task storage with auto-polling
tasks_db = {}
polling_threads = {}

def debug_log(message):
    """Debug logging"""
    if DEBUG:
        timestamp = datetime.now().strftime("%H:%M:%S")
        print(f"[{timestamp}] {message}")

def make_api_request(method, url, **kwargs):
    """Make API request with detailed error logging"""
    try:
        debug_log(f"Making {method} request to {url}")
        response = requests.request(method, url, **kwargs)
        debug_log(f"Response status: {response.status_code}")
        
        if DEBUG and len(response.content) < 10000:  # Don't log huge responses
            debug_log(f"Response: {response.text[:500]}")
        
        return response
    except Exception as e:
        debug_log(f"API request error: {str(e)}")
        raise

def poll_task_status(task_id: str, max_attempts=120, interval=3):
    """Background thread to automatically poll task status"""
    if task_id not in tasks_db:
        debug_log(f"Task {task_id} not found in DB")
        return
    
    task = tasks_db[task_id]
    api_task_id = task["api_task_id"]
    
    debug_log(f"Starting auto-polling for task {task_id} (API ID: {api_task_id})")
    
    for attempt in range(max_attempts):
        try:
            debug_log(f"Poll attempt {attempt + 1} for task {task_id}")
            
            # Poll the API
            url = f"https://api.sunoapi.org/api/v1/lyrics/details?taskId={api_task_id}"
            headers = {"Authorization": f"Bearer {SUNO_KEY}"}
            
            response = make_api_request("GET", url, headers=headers, timeout=30)
            data = response.json()
            
            if response.status_code == 200 and data.get("code") == 200:
                task_data = data["data"]
                status = task_data.get("status", "unknown")
                
                debug_log(f"Task {task_id} status: {status}")
                
                # Update task status
                tasks_db[task_id]["status"] = status
                tasks_db[task_id]["last_checked"] = datetime.now().isoformat()
                
                if status == "completed" and "data" in task_data:
                    # Task completed successfully
                    lyrics_data = task_data["data"]
                    tasks_db[task_id]["result"] = lyrics_data
                    tasks_db[task_id]["completed_at"] = datetime.now().isoformat()
                    debug_log(f"✅ Task {task_id} completed successfully!")
                    break
                
                elif status == "failed":
                    error_msg = task_data.get("error", "Unknown error")
                    tasks_db[task_id]["error"] = error_msg
                    tasks_db[task_id]["completed_at"] = datetime.now().isoformat()
                    debug_log(f"❌ Task {task_id} failed: {error_msg}")
                    break
                elif status == "processing":
                    debug_log(f"🔄 Task {task_id} is processing...")
                else:
                    debug_log(f"📊 Task {task_id} has unknown status: {status}")
            
            else:
                error_code = data.get("code", "unknown")
                error_msg = data.get("msg", "No error message")
                debug_log(f"❌ API error for task {task_id}: Code {error_code}, {error_msg}")
            
            # Update attempts
            tasks_db[task_id]["poll_attempts"] = attempt + 1
            
        except Exception as e:
            debug_log(f"⚠️ Polling error for task {task_id}: {str(e)}")
        
        # Wait before next poll
        time.sleep(interval)
    
    debug_log(f"Finished polling for task {task_id} after {max_attempts} attempts")
    
    # Clean up polling thread
    if task_id in polling_threads:
        del polling_threads[task_id]

def generate_lyrics(prompt: str, auto_poll: bool = True) -> str:
    """Submit lyrics generation task to Suno API"""
    if not SUNO_KEY:
        return "❌ Error: SunoKey environment variable not set. Please add it in Space Settings."
    
    if not prompt or not prompt.strip():
        return "❌ Please enter a lyrics prompt"
    
    # Generate task ID
    task_id = str(uuid.uuid4())[:8]
    
    debug_log(f"Starting generation for task {task_id} with prompt: {prompt[:50]}...")
    
    # Prepare API request
    url = "https://api.sunoapi.org/api/v1/lyrics"
    headers = {
        "Authorization": f"Bearer {SUNO_KEY}",
        "Content-Type": "application/json"
    }
    
    # Suno requires a callback URL
    dummy_callback = "https://dummy.callback.url/not-used"
    
    payload = {
        "prompt": prompt,
        "callBackUrl": dummy_callback
    }
    
    try:
        debug_log(f"Submitting to Suno API: {json.dumps(payload, indent=2)}")
        
        # Submit task
        response = make_api_request("POST", url, headers=headers, json=payload, timeout=30)
        data = response.json()
        
        debug_log(f"Submission response: {json.dumps(data, indent=2)[:500]}...")
        
        if response.status_code == 200 and data.get("code") == 200:
            api_task_id = data["data"]["taskId"]
            
            debug_log(f"✅ Submission successful! Task ID: {task_id}, API Task ID: {api_task_id}")
            
            # Store task information
            tasks_db[task_id] = {
                "id": task_id,
                "api_task_id": api_task_id,
                "prompt": prompt,
                "status": "submitted",
                "result": None,
                "error": None,
                "created_at": datetime.now().isoformat(),
                "last_checked": None,
                "poll_attempts": 0,
                "completed_at": None,
                "auto_poll": auto_poll,
                "raw_response": data  # Store for debugging
            }
            
            # Start auto-polling if enabled
            if auto_poll:
                poll_thread = threading.Thread(
                    target=poll_task_status,
                    args=(task_id,),
                    daemon=True
                )
                polling_threads[task_id] = poll_thread
                poll_thread.start()
                debug_log(f"🚀 Started auto-polling thread for task {task_id}")
            
            return f"""✅ **Task Submitted Successfully!**

**Your Task ID:** `{task_id}`
**API Task ID:** `{api_task_id}`

📝 **Prompt:** {prompt[:100]}{'...' if len(prompt) > 100 else ''}

🔄 **Auto-polling enabled** - Results will appear automatically!

⏳ **Estimated time:** Usually 10-60 seconds

📊 **Check Status tab** to monitor progress

💡 **Save your Task ID:** `{task_id}`"""
        
        else:
            error_msg = data.get("msg", f"HTTP {response.status_code}")
            debug_log(f"❌ Submission failed: {error_msg}")
            
            return f"""❌ **Submission Failed**

**Error:** {error_msg}

**Response:** {json.dumps(data, indent=2)[:500]}

💡 **Possible solutions:**
• Check if your SunoKey is valid
• Ensure API has available credits
• Try a different prompt
• Wait a few minutes and retry"""
    
    except requests.exceptions.Timeout:
        debug_log("❌ Request timeout")
        return "❌ Error: Request timeout - Suno API is not responding"
    except requests.exceptions.ConnectionError:
        debug_log("❌ Connection error")
        return "❌ Error: Connection failed - Check your internet connection"
    except Exception as e:
        debug_log(f"❌ Unexpected error: {str(e)}")
        return f"❌ Error: {str(e)}"

def check_task_status(task_id: str, force_check: bool = False) -> str:
    """Check the status of a task with detailed debugging"""
    if not task_id or not task_id.strip():
        return "❌ Please enter a Task ID"
    
    if task_id not in tasks_db:
        return f"❌ Task ID `{task_id}` not found. Please submit a task first."
    
    task = tasks_db[task_id]
    api_task_id = task.get("api_task_id", "unknown")
    
    debug_log(f"Checking status for task {task_id} (API: {api_task_id})")
    
    # Force an immediate API check if requested
    if force_check:
        try:
            debug_log(f"Forcing API check for task {task_id}")
            url = f"https://api.sunoapi.org/api/v1/lyrics/details?taskId={api_task_id}"
            headers = {"Authorization": f"Bearer {SUNO_KEY}"}
            
            response = make_api_request("GET", url, headers=headers, timeout=30)
            data = response.json()
            
            if response.status_code == 200 and data.get("code") == 200:
                task_data = data["data"]
                status = task_data.get("status", "unknown")
                debug_log(f"Forced check result: {status}")
                
                tasks_db[task_id]["status"] = status
                tasks_db[task_id]["last_checked"] = datetime.now().isoformat()
                
                if status == "completed" and "data" in task_data:
                    lyrics_data = task_data["data"]
                    tasks_db[task_id]["result"] = lyrics_data
                    tasks_db[task_id]["completed_at"] = datetime.now().isoformat()
                elif status == "failed":
                    error_msg = task_data.get("error", "Unknown error")
                    tasks_db[task_id]["error"] = error_msg
                    tasks_db[task_id]["completed_at"] = datetime.now().isoformat()
        
        except Exception as e:
            debug_log(f"Force check error: {str(e)}")
    
    status = task.get("status", "unknown")
    created_time = datetime.fromisoformat(task["created_at"])
    elapsed = int((datetime.now() - created_time).total_seconds())
    
    debug_log(f"Task {task_id} - Status: {status}, Elapsed: {elapsed}s")
    
    # Display based on status
    if status == "completed" and task.get("result"):
        debug_log(f"Task {task_id} has completed results")
        return format_lyrics_output(task["result"], task_id, elapsed)
    
    elif status == "failed":
        error_msg = task.get("error", "Unknown error")
        debug_log(f"Task {task_id} failed: {error_msg}")
        return f"""❌ **Task Failed**

**Task ID:** `{task_id}`
**API Task ID:** `{api_task_id}`
**Error:** {error_msg}
**Elapsed time:** {elapsed} seconds

💡 Please try generating again with a different prompt."""
    
    else:
        # Still processing or unknown status
        attempts = task.get("poll_attempts", 0)
        last_checked = task.get("last_checked")
        last_checked_str = ""
        
        if last_checked:
            last_time = datetime.fromisoformat(last_checked)
            last_checked_str = f"\n**Last checked:** {last_time.strftime('%H:%M:%S')}"
        
        debug_info = ""
        if DEBUG and "raw_response" in task:
            debug_info = f"\n\n**Debug Info:**\n```json\n{json.dumps(task.get('raw_response', {}), indent=2)[:300]}...\n```"
        
        return f"""⏳ **Task Processing...**

**Task ID:** `{task_id}`
**API Task ID:** `{api_task_id}`
**Status:** {status}
**Elapsed:** {elapsed} seconds
**Poll attempts:** {attempts}{last_checked_str}

⏰ **Status Guide:**
- **submitted:** Task accepted by API
- **processing:** AI is generating lyrics
- **completed:** Ready! Check auto-refresh
- **failed:** Error occurred

🔄 **Auto-polling:** {'✅ Active' if task.get('auto_poll') else '❌ Disabled'}

💡 **What to do:**
1. Wait 30-60 seconds for processing
2. Click "Force Check" for immediate update
3. Results appear automatically when ready{debug_info}"""

def force_check_task(task_id: str):
    """Force an immediate API check"""
    return check_task_status(task_id, force_check=True)

def format_lyrics_output(lyrics_data, task_id, elapsed_time):
    """Format the lyrics for display"""
    if not lyrics_data:
        return "✅ Task completed but no lyrics data received"
    
    output_lines = [
        f"# 🎵 Generated Lyrics (Task: {task_id})",
        "",
        f"**Generated in:** {elapsed_time} seconds",
        f"**Completed at:** {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}",
        ""
    ]
    
    for i, item in enumerate(lyrics_data, 1):
        title = item.get('title', f'Variant {i}')
        lyrics = item.get('text', 'No lyrics generated')
        
        output_lines.append(f"## Variant {i}: {title}")
        output_lines.append("```")
        output_lines.append(lyrics)
        output_lines.append("```")
        output_lines.append("---")
        output_lines.append("")
    
    output_lines.append("### 🎉 All done!")
    output_lines.append("You can generate more lyrics or try different prompts.")
    
    return "\n".join(output_lines)

def list_all_tasks():
    """List all submitted tasks with detailed info"""
    if not tasks_db:
        return "📭 No tasks found. Generate some lyrics first!"
    
    output_lines = ["# 📋 All Submitted Tasks", ""]
    
    for task_id, task in sorted(tasks_db.items(), key=lambda x: x[1]["created_at"], reverse=True):
        status = task.get("status", "unknown")
        api_task_id = task.get("api_task_id", "unknown")
        prompt_preview = task.get("prompt", "")[:50]
        created = task.get("created_at", "")[:19]
        
        # Status icon and color
        if status == "completed":
            icon = "✅"
            color = "green"
        elif status in ["failed", "error"]:
            icon = "❌"
            color = "red"
        else:
            icon = "⏳"
            color = "orange"
        
        # Age calculation
        created_time = datetime.fromisoformat(task["created_at"])
        age_seconds = int((datetime.now() - created_time).total_seconds())
        
        output_lines.append(f"<span style='color:{color}'>{icon} **{task_id}** - {status} ({age_seconds}s)</span>")
        output_lines.append(f"   API ID: `{api_task_id}`")
        output_lines.append(f"   Prompt: {prompt_preview}...")
        output_lines.append(f"   Created: {created}")
        output_lines.append(f"   Poll attempts: {task.get('poll_attempts', 0)}")
        
        if task.get("last_checked"):
            last_checked = task["last_checked"][:19]
            output_lines.append(f"   Last checked: {last_checked}")
        
        if task.get("completed_at"):
            completed = task["completed_at"][:19]
            output_lines.append(f"   Completed: {completed}")
        
        output_lines.append("")
    
    # Add summary
    completed = sum(1 for t in tasks_db.values() if t.get("status") == "completed")
    processing = sum(1 for t in tasks_db.values() if t.get("status") not in ["completed", "failed", "error"])
    failed = sum(1 for t in tasks_db.values() if t.get("status") in ["failed", "error"])
    total = len(tasks_db)
    
    output_lines.append(f"**Summary:** {completed} ✅, {processing} ⏳, {failed} ❌, {total} total")
    output_lines.append(f"**Active poll threads:** {len(polling_threads)}")
    
    return "\n".join(output_lines)

def get_diagnostic_info():
    """Get diagnostic information about the system"""
    info_lines = ["# 🩺 Diagnostic Information", ""]
    
    # API Key status
    api_status = "✅ Configured" if SUNO_KEY else "❌ NOT SET"
    api_preview = SUNO_KEY[:10] + "..." if SUNO_KEY and len(SUNO_KEY) > 10 else SUNO_KEY or "None"
    info_lines.append(f"**SunoKey:** {api_status} ({api_preview})")
    
    # Tasks summary
    info_lines.append(f"\n**Tasks in memory:** {len(tasks_db)}")
    info_lines.append(f"**Active poll threads:** {len(polling_threads)}")
    
    # Recent tasks
    if tasks_db:
        info_lines.append("\n**Recent Tasks:**")
        for task_id, task in sorted(tasks_db.items(), key=lambda x: x[1]["created_at"], reverse=True)[:5]:
            status = task.get("status", "unknown")
            age = int((datetime.now() - datetime.fromisoformat(task["created_at"])).total_seconds())
            info_lines.append(f"- `{task_id}`: {status} ({age}s ago)")
    
    # Test API connection
    info_lines.append("\n**API Connection Test:**")
    try:
        test_response = requests.get("https://api.sunoapi.org", timeout=5)
        info_lines.append(f"✅ Suno API reachable (HTTP {test_response.status_code})")
    except Exception as e:
        info_lines.append(f"❌ Cannot reach Suno API: {str(e)}")
    
    # Debug info
    info_lines.append(f"\n**Debug Mode:** {'✅ Enabled' if DEBUG else '❌ Disabled'}")
    info_lines.append(f"**Current Time:** {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    
    return "\n".join(info_lines)

# Create the Gradio interface
with gr.Blocks(title="Suno Lyrics Generator", theme=gr.themes.Soft()) as app:
    gr.Markdown("# 🎵 Suno AI Lyrics Generator")
    gr.Markdown("Generate song lyrics using Suno's AI API")
    
    # Store the current task ID in a hidden state
    current_task_id = gr.State("")
    
    with gr.Tabs():
        # Tab 1: Generate Lyrics
        with gr.TabItem("✨ Generate"):
            with gr.Row():
                with gr.Column(scale=2):
                    gr.Markdown("### Enter your lyrics idea")
                    prompt_input = gr.Textbox(
                        label="Lyrics Prompt",
                        placeholder="Example: A romantic ballad about stargazing on a summer night...",
                        lines=3
                    )
                    
                    auto_poll_checkbox = gr.Checkbox(
                        label="🔁 Enable auto-polling",
                        value=True,
                        info="Automatically check for results (recommended)"
                    )
                    
                    submit_btn = gr.Button("🚀 Generate Lyrics", variant="primary", size="lg")
                    
                    gr.Markdown("### ⚠️ Current Issue:")
                    gr.Markdown("""
                    Tasks are getting stuck in "submitted" status.
                    This version includes **debugging tools** to diagnose the issue.
                    
                    **What to try:**
                    1. Submit a simple prompt
                    2. Check "Diagnostic" tab
                    3. Use "Force Check" button
                    4. Monitor debug logs
                    """)
                
                with gr.Column(scale=3):
                    output_area = gr.Markdown(
                        label="Result",
                        value="Your task submission result will appear here..."
                    )
            
            def handle_generation(prompt, auto_poll):
                result = generate_lyrics(prompt, auto_poll)
                # Extract task ID from result
                task_id = ""
                if "Task ID:" in result:
                    for line in result.split('\n'):
                        if 'Task ID:' in line:
                            parts = line.split('`')
                            if len(parts) > 1:
                                task_id = parts[1]
                                break
                return result, task_id
            
            submit_btn.click(
                fn=handle_generation,
                inputs=[prompt_input, auto_poll_checkbox],
                outputs=[output_area, current_task_id]
            )
        
        # Tab 2: Check Status
        with gr.TabItem("🔍 Check Status"):
            with gr.Row():
                with gr.Column():
                    gr.Markdown("### Enter your Task ID")
                    task_id_input = gr.Textbox(
                        label="Task ID",
                        placeholder="Paste your Task ID here (e.g., 0f015fcb)",
                        scale=1
                    )
                    
                    # Auto-populate if we have a current task ID
                    def update_task_id_input(current_id):
                        return current_id if current_id else ""
                    
                    current_task_id.change(
                        fn=update_task_id_input,
                        inputs=current_task_id,
                        outputs=task_id_input
                    )
                    
                    with gr.Row():
                        check_btn = gr.Button("🔍 Check Status", variant="primary")
                        force_check_btn = gr.Button("⚡ Force Check", variant="secondary")
                        auto_refresh_btn = gr.Button("🔄 Auto-refresh", variant="secondary")
                    
                    gr.Markdown("---")
                    refresh_all_btn = gr.Button("📋 List All Tasks")
                    tasks_list = gr.Markdown(label="All Tasks")
                
                with gr.Column():
                    status_output = gr.Markdown(
                        label="Status",
                        value="Enter a Task ID above and click Check Status"
                    )
            
            # Regular check
            check_btn.click(
                fn=check_task_status,
                inputs=task_id_input,
                outputs=status_output
            )
            
            # Force check (immediate API call)
            force_check_btn.click(
                fn=force_check_task,
                inputs=task_id_input,
                outputs=status_output
            )
            
            # List all tasks
            refresh_all_btn.click(
                fn=list_all_tasks,
                inputs=None,
                outputs=tasks_list
            )
        
        # Tab 3: Diagnostic
        with gr.TabItem("🩺 Diagnostic"):
            gr.Markdown("# System Diagnostics")
            gr.Markdown("Use this tab to diagnose why tasks are getting stuck")
            
            with gr.Row():
                with gr.Column():
                    diagnostic_btn = gr.Button("🔄 Refresh Diagnostics", variant="primary")
                    diagnostic_output = gr.Markdown(label="Diagnostic Info")
                
                with gr.Column():
                    gr.Markdown("### 🚨 Common Issues:")
                    gr.Markdown("""
                    **1. API Key Issues:**
                    - Invalid or expired SunoKey
                    - No API credits remaining
                    - Incorrect environment variable name
                    
                    **2. API Response Issues:**
                    - Suno API returning errors
                    - Tasks stuck in queue
                    - Rate limiting
                    
                    **3. Network Issues:**
                    - Cannot reach api.sunoapi.org
                    - Timeout errors
                    - Connection refused
                    
                    **4. Task Processing:**
                    - Suno AI taking longer than expected
                    - Tasks stuck in "submitted" state
                    - Server-side delays
                    """)
            
            diagnostic_btn.click(
                fn=get_diagnostic_info,
                inputs=None,
                outputs=diagnostic_output
            )
            
            # Auto-refresh diagnostics every 10 seconds
            diagnostic_btn.click(
                fn=lambda: time.sleep(10),
                inputs=None,
                outputs=None
            ).then(
                fn=get_diagnostic_info,
                inputs=None,
                outputs=diagnostic_output
            )
        
        # Tab 4: Help
        with gr.TabItem("ℹ️ Help"):
            gr.Markdown("# Help & Troubleshooting")
            
            with gr.Row():
                with gr.Column():
                    gr.Markdown("### 🐛 Debugging Stuck Tasks")
                    gr.Markdown("""
                    **If tasks are stuck in "submitted":**
                    
                    1. **Check Diagnostic Tab:**
                       - Verify API key is set
                       - Test API connection
                       - View recent task status
                    
                    2. **Use Force Check:**
                       - Makes immediate API call
                       - Bypasses cached status
                       - Shows raw API response
                    
                    3. **Monitor Debug Logs:**
                       - Check Space logs (bottom of page)
                       - Look for API errors
                       - Note timeout messages
                    
                    4. **Try Simple Test:**
                       - Use a very simple prompt
                       - Disable auto-polling
                       - Check after 60 seconds
                    """)
                
                with gr.Column():
                    gr.Markdown("### 📞 Support")
                    gr.Markdown("""
                    **If issues persist:**
                    
                    1. **Check SunoKey:**
                       - Ensure it's valid
                       - Check credit balance
                       - Try in Suno's own interface
                    
                    2. **API Status:**
                       - Suno API may be down
                       - Check Suno status page
                       - Wait and try later
                    
                    3. **Contact Support:**
                       - Suno API support
                       - Provide your Task IDs
                       - Share debug logs
                    
                    4. **Alternative:**
                       - Try a different prompt
                       - Wait 5 minutes
                       - Restart the Space
                    """)

# Launch the app
if __name__ == "__main__":
    print("=" * 60)
    print("🚀 Starting Suno Lyrics Generator - DEBUG MODE")
    print("=" * 60)
    print(f"🔑 SunoKey: {'✅ Configured' if SUNO_KEY else '❌ NOT SET'}")
    if SUNO_KEY:
        print(f"🔑 Preview: {SUNO_KEY[:10]}...")
    print(f"🐛 Debug Mode: {'✅ Enabled' if DEBUG else '❌ Disabled'}")
    print(f"📊 Pre-existing tasks: {len(tasks_db)}")
    print("=" * 60)
    
    app.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        debug=False,
        show_error=True
    )