import gradio as gr
import os
import tempfile
import shutil
import requests
import json
import uuid
import time
from datetime import datetime
from dotenv import load_dotenv
from urllib.parse import urlparse, parse_qs

# Load environment variables
load_dotenv()

# Configuration
SUNO_API_KEY = os.environ.get("SunoKey", "")
FIXED_CALLBACK_URL = "https://1hit.no/cover/cb.php"
SUNO_API_URL = "https://api.sunoapi.org/api/v1/generate/upload-cover"

# Create uploads directory
UPLOADS_FOLDER = "uploads"
os.makedirs(UPLOADS_FOLDER, exist_ok=True)

# ============================================
# PAGE 1: FILE UPLOAD
# ============================================

def upload_file(file_obj, custom_name):
    """Simple file upload - returns file for download"""
    if not file_obj:
        return None, "❌ Please upload a file first"
    
    try:
        # Handle Gradio 6+ file object
        if isinstance(file_obj, dict):
            file_path = file_obj["path"]
            original_name = file_obj["name"]
        else:
            file_path = file_obj.name
            original_name = os.path.basename(file_path)
        
        # Create temp directory
        temp_dir = tempfile.mkdtemp()
        
        # Determine filename
        if custom_name and custom_name.strip():
            ext = os.path.splitext(original_name)[1]
            base_name = custom_name.strip()
            if not base_name.endswith(ext):
                final_name = base_name + ext
            else:
                final_name = base_name
        else:
            final_name = original_name
        
        # Copy file
        final_path = os.path.join(temp_dir, final_name)
        shutil.copy2(file_path, final_path)
        
        return final_path, f"✅ Ready: {final_name}"
    
    except Exception as e:
        return None, f"❌ Error: {str(e)}"

def get_file_info(file_obj):
    """Display file information"""
    if not file_obj:
        return "📁 No file selected"
    
    try:
        if isinstance(file_obj, dict):
            file_path = file_obj["path"]
            file_name = file_obj["name"]
        else:
            file_path = file_obj.name
            file_name = os.path.basename(file_path)
        
        size = os.path.getsize(file_path)
        return f"📄 **{file_name}**\n• Size: {size/1024:.1f} KB\n• Type: {os.path.splitext(file_name)[1]}"
    except:
        return "📁 File selected"

# ============================================
# PAGE 2: COVER CREATION
# ============================================

def generate_cover(
    prompt,
    title,
    style,
    upload_url_type,
    custom_upload_url,
    uploaded_file,
    instrumental=True,
    model="V4_5ALL",
    persona_id="",
    negative_tags="",
    vocal_gender="m",
    style_weight=0.65,
    weirdness_constraint=0.65,
    audio_weight=0.65,
    custom_mode=True
):
    """Generate cover using Suno API"""
    
    # Check API key
    if not SUNO_API_KEY:
        return "❌ Suno API key not found. Please set 'SunoKey' environment variable."
    
    # Determine upload URL
    if upload_url_type == "uploaded":
        # Use uploaded file - in real implementation, you'd upload to storage and get URL
        if not uploaded_file:
            return "❌ Please upload a file first or select another URL type"
        
        # For demo - simulate URL from uploaded file
        if isinstance(uploaded_file, dict):
            file_name = uploaded_file.get("name", "uploaded_file.mp3")
        else:
            file_name = getattr(uploaded_file, "name", "uploaded_file.mp3")
        
        upload_url = f"https://storage.temp.example.com/uploads/{int(time.time())}_{file_name}"
        
    elif upload_url_type == "custom":
        upload_url = custom_upload_url.strip()
        if not upload_url:
            return "❌ Please provide a custom upload URL"
    else:  # auto
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        unique_id = str(uuid.uuid4())[:8]
        upload_url = f"https://storage.temp.example.com/uploads/{timestamp}_{unique_id}.mp3"
    
    # Prepare payload
    payload = {
        "uploadUrl": upload_url,
        "customMode": custom_mode,
        "instrumental": instrumental,
        "model": model,
        "callBackUrl": FIXED_CALLBACK_URL,
        "prompt": prompt,
        "style": style,
        "title": title,
        "personaId": persona_id,
        "negativeTags": negative_tags,
        "vocalGender": vocal_gender,
        "styleWeight": style_weight,
        "weirdnessConstraint": weirdness_constraint,
        "audioWeight": audio_weight
    }
    
    # Remove empty fields
    payload = {k: v for k, v in payload.items() if v not in ["", None]}
    
    # Headers
    headers = {
        "Authorization": f"Bearer {SUNO_API_KEY}",
        "Content-Type": "application/json"
    }
    
    try:
        # Make API request
        response = requests.post(SUNO_API_URL, json=payload, headers=headers)
        
        if response.status_code == 200:
            result = response.json()
            if result.get("code") == 200:
                task_id = result.get("data", {}).get("taskId", "Unknown")
                return f"""✅ **Generation Started!**

**Task ID:** `{task_id}`

**Upload URL:** {upload_url}
**Callback URL:** {FIXED_CALLBACK_URL}

📋 **Full Response:**
```json
{json.dumps(result, indent=2)}
```"""
            else:
                return f"""❌ **API Error:** {result.get('msg', 'Unknown error')}

📋 **Response:**
```json
{json.dumps(result, indent=2)}
```"""
        else:
            return f"""❌ **HTTP Error {response.status_code}**

{response.text}"""
            
    except Exception as e:
        return f"❌ **Error:** {str(e)}"

# ============================================
# PAGE 3: CHECK TASK - ORIGINAL CODE
# ============================================

def get_task_info(task_id):
    """Manually check any Suno task status - ORIGINAL CODE"""
    if not task_id:
        return "❌ Please enter a Task ID ❌"
    
    try:
        resp = requests.get(
            "https://api.sunoapi.org/api/v1/generate/record-info",
            headers={"Authorization": f"Bearer {SUNO_API_KEY}"},
            params={"taskId": task_id},
            timeout=30
        )
        
        if resp.status_code != 200:
            return f"❌ HTTP Error {resp.status_code}\n\n{resp.text}"
        
        data = resp.json()
        
        # Format the response for display
        output = f"## 🔍 Task Status: `{task_id}`\n\n"
        
        if data.get("code") == 200:
            task_data = data.get("data", {})
            status = task_data.get("status", "UNKNOWN")
            
            output += f"**Status:** {status}\n"
            output += f"**Task ID:** `{task_data.get('taskId', 'N/A')}`\n"
            output += f"**Music ID:** `{task_data.get('musicId', 'N/A')}`\n"
            output += f"**Created:** {task_data.get('createTime', 'N/A')}\n"
            
            if status == "SUCCESS":
                response_data = task_data.get("response", {})
                
                # Try to parse response (could be string or dict)
                if isinstance(response_data, str):
                    try:
                        response_data = json.loads(response_data)
                    except:
                        output += f"\n**Raw Response:**\n```\n{response_data}\n```\n"
                        response_data = {}
                
                # Check for song data
                songs = []
                if isinstance(response_data, dict):
                    songs = response_data.get("sunoData", [])
                    if not songs:
                        songs = response_data.get("data", [])
                elif isinstance(response_data, list):
                    songs = response_data
                
                if songs:
                    output += f"\n## 🎵 Generated Songs ({len(songs)})\n\n"
                    
                    for i, song in enumerate(songs, 1):
                        if isinstance(song, dict):
                            output += f"### Song {i}\n"
                            output += f"**Title:** {song.get('title', 'Untitled')}\n"
                            output += f"**ID:** `{song.get('id', 'N/A')}`\n"
                            
                            # Audio URLs
                            audio_url = song.get('audioUrl') or song.get('audio_url')
                            stream_url = song.get('streamUrl') or song.get('stream_url')
                            download_url = song.get('downloadUrl') or song.get('download_url')
                            
                            if audio_url:
                                output += f"**Audio:** [Play]({audio_url}) | [Download]({audio_url})\n"
                            elif stream_url:
                                output += f"**Stream:** [Play]({stream_url})\n"
                            
                            if download_url:
                                output += f"**Download:** [MP3]({download_url})\n"
                            
                            # Audio player
                            play_url = audio_url or stream_url
                            if play_url:
                                output += f"""\n<audio controls style="width: 100%; margin: 10px 0;">
                                    <source src="{play_url}" type="audio/mpeg">
                                    Your browser does not support audio.
                                </audio>\n"""
                            
                            output += f"**Prompt:** {song.get('prompt', 'N/A')[:100]}...\n"
                            output += f"**Duration:** {song.get('duration', 'N/A')}s\n"
                            output += f"**Created:** {song.get('createTime', 'N/A')}\n\n"
                            output += "---\n\n"
                else:
                    output += "\n**No song data found in response.**\n"
                    
            elif status == "FAILED":
                error_msg = task_data.get("errorMessage", "Unknown error")
                output += f"\n**Error:** {error_msg}\n"
            
            elif status in ["PENDING", "PROCESSING", "RUNNING"]:
                output += f"\n**Task is still processing...**\n"
                output += f"Check again in 30 seconds.\n"
            
            else:
                output += f"\n**Unknown status:** {status}\n"
        
        else:
            output += f"**API Error:** {data.get('msg', 'Unknown')}\n"
        
        # Show raw JSON for debugging
        output += "\n## 📋 Raw Response\n"
        output += f"```json\n{json.dumps(data, indent=2)}\n```"
        
        return output
        
    except Exception as e:
        return f"❌ Error checking task: {str(e)}"

# ============================================
# URL PARAMETER HANDLING
# ============================================

def parse_url_params(request: gr.Request):
    """Parse taskid from URL parameters"""
    task_id = None
    if request:
        try:
            query_params = parse_qs(urlparse(request.request.url).query)
            if 'taskid' in query_params:
                task_id = query_params['taskid'][0]
                # Remove any whitespace
                task_id = task_id.strip()
        except Exception as e:
            print(f"Error parsing URL params: {e}")
    
    return task_id

def on_page_load(request: gr.Request):
    """Handle URL parameters when page loads"""
    task_id = parse_url_params(request)
    
    if task_id:
        # We have a task ID from URL, return it and fetch results
        task_result = get_task_info(task_id)
        return (
            task_id,  # For check_task_id
            task_result,  # For check_output
            gr.Tabs(selected="tab_check"),  # Switch to check tab
            True  # Mark as loaded
        )
    else:
        # No task ID in URL, stay on first tab
        return (
            "",  # Empty check_task_id
            "### Enter a Task ID above\n\nPaste any Suno Task ID to check its current status and results.",  # Default message
            gr.Tabs(selected="tab_upload"),  # Stay on upload tab
            True  # Mark as loaded
        )

# ============================================
# BUILD THE APP
# ============================================

# Custom CSS
css = """
    .status-badge { 
        padding: 5px 10px; 
        border-radius: 20px; 
        font-weight: bold; 
    }
    .api-status-ok { background-color: #d4edda; color: #155724; }
    .api-status-missing { background-color: #f8d7da; color: #721c24; }
    audio { width: 100%; margin: 10px 0; }
"""

# Build interface
with gr.Blocks(title="Suno Cover Creator", theme=gr.themes.Soft(), css=css) as app:
    
    gr.Markdown("""
    <div style="text-align: center; margin-bottom: 20px;">
        <h1>🎵 Suno Cover Creator</h1>
        <p>Upload → Create Cover → Check Status</p>
    </div>
    """)
    
    # API Status Banner
    api_status_color = "✅ API Key: Loaded" if SUNO_API_KEY else "❌ API Key: Not Found"
    api_status_class = "api-status-ok" if SUNO_API_KEY else "api-status-missing"
    
    with gr.Row():
        gr.Markdown(f"""
        <div class="{api_status_class}" style="padding: 10px; border-radius: 5px; text-align: center;">
            {api_status_color}
        </div>
        """)
    
    # State for initial load tracking
    initial_load_done = gr.State(value=False)
    
    # Main tabs
    with gr.Tabs() as tabs:
        # ========== PAGE 1: FILE UPLOAD ==========
        with gr.TabItem("📤 1. Upload File", id="tab_upload"):
            with gr.Row():
                with gr.Column():
                    file_input = gr.File(
                        label="Upload Your Audio File",
                        file_count="single",
                        file_types=["audio", ".mp3", ".wav", ".m4a", ".flac", ".ogg"]
                    )
                    
                    file_info = gr.Markdown("📁 No file selected")
                    
                    custom_filename = gr.Textbox(
                        label="Custom Filename (optional)",
                        placeholder="my-audio-file",
                        info="Extension preserved automatically"
                    )
                    
                    upload_btn = gr.Button("⚡ Prepare File", variant="primary", size="lg")
                    
                    upload_status = gr.Textbox(label="Status", interactive=False)
                    upload_download = gr.File(label="Download File", interactive=False)
        
        # ========== PAGE 2: COVER CREATION ==========
        with gr.TabItem("🎨 2. Create Cover", id="tab_create"):
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("### 🎵 Cover Details")
                    
                    cover_prompt = gr.Textbox(
                        label="Prompt",
                        value="A dramatic orchestral cover with dark undertones",
                        lines=2
                    )
                    
                    cover_title = gr.Textbox(
                        label="Title",
                        value="The Fool's Ascension"
                    )
                    
                    cover_style = gr.Textbox(
                        label="Style",
                        value="Epic Orchestral, Dark Cinematic"
                    )
                    
                    # URL Selection
                    gr.Markdown("### 🔗 Upload URL")
                    
                    url_type = gr.Radio(
                        label="Source",
                        choices=[
                            ("Use uploaded file", "uploaded"),
                            ("Use custom URL", "custom"),
                            ("Auto-generate", "auto")
                        ],
                        value="uploaded",
                        info="Choose where the audio comes from"
                    )
                    
                    custom_url = gr.Textbox(
                        label="Custom URL",
                        placeholder="https://example.com/audio.mp3",
                        visible=False
                    )
                    
                    # Reference to uploaded file from Page 1
                    uploaded_file_ref = gr.State()
                    
                with gr.Column(scale=1):
                    gr.Markdown("### ⚙️ Advanced Settings")
                    
                    cover_model = gr.Dropdown(
                        label="Model",
                        choices=["V4_5ALL", "V5", "V4"],
                        value="V4_5ALL"
                    )
                    
                    cover_instrumental = gr.Checkbox(
                        label="Instrumental",
                        value=True
                    )
                    
                    cover_vocal = gr.Dropdown(
                        label="Vocal Gender",
                        choices=["m", "f", "none"],
                        value="m"
                    )
                    
                    cover_negative = gr.Textbox(
                        label="Negative Tags",
                        value="Distorted, Low Quality",
                        placeholder="What to avoid"
                    )
            
            # Show/hide custom URL
            def toggle_custom_url(choice):
                return gr.update(visible=choice == "custom")
            
            url_type.change(
                toggle_custom_url,
                inputs=url_type,
                outputs=custom_url
            )
            
            # Generate button
            generate_btn = gr.Button("🎬 Generate Cover", variant="primary", size="lg")
            generate_output = gr.Markdown("Ready to generate...")
        
        # ========== PAGE 3: CHECK TASK - ORIGINAL ==========
        with gr.TabItem("🔍 3. Check Any Task", id="tab_check"):
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("### Check Task Status")
                    gr.Markdown("Enter any Suno Task ID to check its status")
                    
                    check_task_id = gr.Textbox(
                        label="Task ID",
                        placeholder="Enter Task ID from generation",
                        info="From Cover Creator or any Suno API request"
                    )
                    
                    check_btn = gr.Button("🔍 Check Status", variant="primary")
                    check_clear_btn = gr.Button("🗑️ Clear", variant="secondary")
                    
                    # URL parameter info
                    gr.Markdown("""
                    **Quick access via URL:**
                    Add `?taskid=YOUR_TASK_ID` to the URL
                    
                    Example:
                    `https://www.1hit.no/cover/logs.php`
                    """)
                
                with gr.Column(scale=2):
                    check_output = gr.Markdown(
                        value="### Enter a Task ID above\n\nPaste any Suno Task ID to check its current status and results."
                    )
    
    # ============================================
    # EVENT HANDLERS
    # ============================================
    
    # Page 1: File Upload
    file_input.change(
        fn=get_file_info,
        inputs=[file_input],
        outputs=file_info
    )
    
    upload_btn.click(
        fn=upload_file,
        inputs=[file_input, custom_filename],
        outputs=[upload_download, upload_status]
    )
    
    # Store uploaded file reference for Page 2
    def store_file_ref(file_obj):
        return file_obj
    
    file_input.change(
        fn=store_file_ref,
        inputs=[file_input],
        outputs=uploaded_file_ref
    )
    
    # Page 2: Generate Cover
    generate_btn.click(
        fn=generate_cover,
        inputs=[
            cover_prompt, cover_title, cover_style,
            url_type, custom_url, uploaded_file_ref,
            cover_instrumental, cover_model,
            gr.State(""),  # persona_id placeholder
            cover_negative, cover_vocal,
            gr.State(0.65), gr.State(0.65), gr.State(0.65),  # weights
            gr.State(True)  # custom_mode
        ],
        outputs=generate_output
    )
    
    # Page 3: Check Task - Original handlers
    def clear_check():
        return "", "### Enter a Task ID above\n\nPaste any Suno Task ID to check its current status and results."
    
    check_clear_btn.click(
        clear_check,
        outputs=[check_task_id, check_output]
    )
    
    check_btn.click(
        get_task_info,
        inputs=[check_task_id],
        outputs=check_output
    )
    
    # Load URL parameters when the app starts
    app.load(
        fn=on_page_load,
        inputs=[],
        outputs=[check_task_id, check_output, tabs, initial_load_done],
        queue=False
    )
    
    # Footer
    gr.Markdown("---")
    gr.Markdown("""
    <div style="text-align: center; padding: 20px;">
        <p>🔗 <b>Quick URL Access:</b> Add <code>?taskid=YOUR_TASK_ID</code> to auto-load task</p>
       
    </div>
    """)

# Launch
if __name__ == "__main__":
    if not SUNO_API_KEY:
        print("⚠️ Warning: Suno API key not found")
        print("Set 'SunoKey' environment variable or add to .env file")
    
    print("🚀 Starting Suno Cover Creator")
    print(f"🔑 SunoKey: {'✅ Set' if SUNO_API_KEY else '❌ Not set'}")
    print("🔗 Use URL parameter: http://localhost:7860/?taskid=YOUR_TASK_ID")
    
    app.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )