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

# 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"
CHECK_TASK_URL = "https://api.sunoapi.org/api/v1/task"  # Adjust if different

# 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
# ============================================

def check_task_status(task_id):
    """Check status of a generation task"""
    if not task_id or not task_id.strip():
        return "❌ Please enter a Task ID"
    
    if not SUNO_API_KEY:
        return "❌ Suno API key not found"
    
    headers = {
        "Authorization": f"Bearer {SUNO_API_KEY}",
        "Content-Type": "application/json"
    }
    
    try:
        # Try different endpoint patterns - adjust based on actual API
        endpoints = [
            f"{CHECK_TASK_URL}/{task_id}",
            f"{CHECK_TASK_URL}?taskId={task_id}",
            f"https://api.sunoapi.org/api/v1/task/{task_id}"
        ]
        
        for url in endpoints:
            try:
                response = requests.get(url, headers=headers)
                if response.status_code == 200:
                    result = response.json()
                    return f"""✅ **Task Status**

📋 **Result:**
```json
{json.dumps(result, indent=2)}
```"""
            except:
                continue
        
        return f"❌ Could not find task with ID: {task_id}"
        
    except Exception as e:
        return f"❌ **Error:** {str(e)}"

def poll_task(task_id, interval=5, max_attempts=12):
    """Poll task until completion"""
    if not task_id or not task_id.strip():
        return "❌ Please enter a Task ID"
    
    status_msg = ""
    for attempt in range(max_attempts):
        status_msg += f"⏳ Polling attempt {attempt + 1}/{max_attempts}...\n"
        
        # Get status
        result = check_task_status(task_id)
        status_msg += f"{result}\n\n"
        
        # Check if completed (adjust based on actual response format)
        if "completed" in result.lower() or "success" in result.lower():
            status_msg += "✅ **Task completed!**"
            break
        
        if attempt < max_attempts - 1:
            time.sleep(interval)
    
    return status_msg

# ============================================
# 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; }
"""

# 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} | Callback URL: <code>{FIXED_CALLBACK_URL}</code>
        </div>
        """)
    
    # Main tabs
    with gr.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", "V3"],
                        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 ==========
        with gr.TabItem("🔍 3. Check Task", id="tab_check"):
            with gr.Row():
                with gr.Column():
                    gr.Markdown("### Check Single Task")
                    
                    task_id_input = gr.Textbox(
                        label="Task ID",
                        placeholder="Enter task ID from generation response",
                        lines=1
                    )
                    
                    check_btn = gr.Button("🔍 Check Status", variant="primary")
                    check_output = gr.Markdown("Enter a Task ID to check...")
                
                with gr.Column():
                    gr.Markdown("### Poll Until Complete")
                    
                    poll_task_id = gr.Textbox(
                        label="Task ID",
                        placeholder="Enter task ID to poll"
                    )
                    
                    poll_interval = gr.Slider(
                        label="Poll Interval (seconds)",
                        minimum=2,
                        maximum=10,
                        value=5,
                        step=1
                    )
                    
                    poll_attempts = gr.Slider(
                        label="Max Attempts",
                        minimum=5,
                        maximum=30,
                        value=12,
                        step=1
                    )
                    
                    poll_btn = gr.Button("⏳ Poll Until Complete", variant="secondary")
                    poll_output = gr.Markdown("Enter a Task ID to poll...")
    
    # ============================================
    # 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
    check_btn.click(
        fn=check_task_status,
        inputs=[task_id_input],
        outputs=check_output
    )
    
    def poll_wrapper(task_id, interval, attempts):
        if not task_id:
            return "❌ Please enter a Task ID"
        
        status_msg = ""
        for attempt in range(int(attempts)):
            status_msg += f"⏳ Attempt {attempt + 1}/{int(attempts)}...\n"
            result = check_task_status(task_id)
            status_msg += f"{result}\n\n"
            
            # Check if completed (simplified)
            if "completed" in result.lower() or "success" in result.lower():
                status_msg += "✅ **Task completed!**"
                break
            
            if attempt < int(attempts) - 1:
                time.sleep(int(interval))
        
        return status_msg
    
    poll_btn.click(
        fn=poll_wrapper,
        inputs=[poll_task_id, poll_interval, poll_attempts],
        outputs=poll_output
    )
    
    # Footer
    gr.Markdown("---")
    gr.Markdown("💡 **Tips:** Upload your audio in Page 1, create cover in Page 2, then check status in Page 3")

# 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")
    
    app.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )