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"
LOGS_URL = "https://www.1hit.no/cover/compu.php"
BACK_URL = "https://www.1hit.no/cover/ib.php"

# 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:
        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)
        
        temp_dir = tempfile.mkdtemp()
        
        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
        
        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 WITH PREVIEW
# ============================================

def get_file_path(file_obj):
    """Extract file path from uploaded file object"""
    if not file_obj:
        return None
    if isinstance(file_obj, dict):
        return file_obj.get("path", "")
    return getattr(file_obj, "name", "")

def generate_cover(
    prompt,
    title,
    style,
    use_custom_url,
    custom_url_input,
    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"""
    
    if not SUNO_API_KEY:
        return "❌ Suno API key not found"
    
    # Determine which URL to use
    if use_custom_url:
        upload_url = custom_url_input.strip()
        if not upload_url:
            return "❌ Please enter a custom URL"
        url_source = f"Custom URL: {upload_url}"
    else:
        if not uploaded_file:
            return "❌ Please upload a file first"
        #upload_url = get_file_path(uploaded_file)
        #upload_url = f"/file={get_file_path(uploaded_file)}"
        upload_url = f"https://1hit-upload-cover.hf.space/gradio_api/file={get_file_path(uploaded_file)}"
        if not upload_url:
            return "❌ Could not get file path from uploaded file"
        
        url_source = f"Uploaded file: {upload_url}"
    
    # 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
    }
    
    payload = {k: v for k, v in payload.items() if v not in ["", None]}
    
    headers = {
        "Authorization": f"Bearer {SUNO_API_KEY}",
        "Content-Type": "application/json"
    }
    
    try:
        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}`

**URL Source:** {url_source}

📋 **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)}"

def update_preview(use_custom_url, custom_url_input, uploaded_file):
    """Update audio preview based on selected source"""
    if use_custom_url:
        url = custom_url_input.strip()
        if url:
            return gr.update(value=url, visible=True, label="Preview URL")
        else:
            return gr.update(value=None, visible=False, label="Preview URL")
    else:
        if uploaded_file:
            file_path = get_file_path(uploaded_file)
            if file_path:
                return gr.update(value=file_path, visible=True, label="Preview from uploaded file")
    return gr.update(value=None, visible=False, label="Preview URL")

# ============================================
# PAGE 3: CHECK TASK
# ============================================

def get_task_info(task_id):
    """Check Suno task status"""
    if not task_id:
        return "❌ Please enter a Task ID ❌"
    
    if not SUNO_API_KEY:
        return "❌ Suno API key not configured"
    
    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 == 401:
            return "❌ Authentication failed - Invalid API key"
        elif resp.status_code == 404:
            return f"❌ Task ID '{task_id}' not found"
        elif resp.status_code != 200:
            return f"❌ HTTP Error {resp.status_code}\n\n{resp.text}"
        
        data = resp.json()
        
        # Format the response
        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", {})
                
                if isinstance(response_data, str):
                    try:
                        response_data = json.loads(response_data)
                    except:
                        response_data = {}
                
                songs = []
                if isinstance(response_data, dict):
                    songs = response_data.get("sunoData", []) or 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_url = song.get('audioUrl') or song.get('audio_url')
                            stream_url = song.get('streamUrl') or song.get('stream_url')
                            
                            if audio_url:
                                output += f"**Audio:** [Play]({audio_url}) | [Download]({audio_url})\n"
                                output += f"""\n<audio controls style="width: 100%; margin: 10px 0;">
                                    <source src="{audio_url}" type="audio/mpeg">
                                    Your browser does not support audio.
                                </audio>\n"""
                            elif stream_url:
                                output += f"**Stream:** [Play]({stream_url})\n"
                                output += f"""\n<audio controls style="width: 100%; margin: 10px 0;">
                                    <source src="{stream_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\n"
                            output += "---\n\n"
                else:
                    output += "\n**No song data found in response.**\n"
        else:
            output += f"**API Error:** {data.get('msg', 'Unknown')}\n"
        
        output += f"\n## 📋 Raw Response\n```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"""
    if request:
        try:
            query_params = parse_qs(urlparse(request.request.url).query)
            if 'taskid' in query_params:
                return query_params['taskid'][0].strip()
        except Exception as e:
            print(f"Error parsing URL params: {e}")
    return None

def on_page_load(request: gr.Request):
    """Handle URL parameters when page loads"""
    task_id = parse_url_params(request)
    
    if task_id:
        task_result = get_task_info(task_id)
        return (
            task_id,
            task_result,
            gr.Tabs(selected="tab_check"),
            True
        )
    else:
        return (
            "",
            "### Enter a Task ID above\n\nPaste any Suno Task ID to check its current status and results.",
            gr.Tabs(selected="tab_upload"),
            True
        )

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

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; }
    .logs-iframe { width: 100%; height: 600px; border: 1px solid #ddd; border-radius: 5px; }
    .preview-box { 
        border: 1px solid #ddd; 
        border-radius: 5px; 
        padding: 15px; 
        margin: 10px 0;
        background-color: #f9f9f9;
    }
"""

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 → Preview Audio → Create Cover → Check Status → View Logs</p>
        <p>Friendly advice: Dont cover originals, cover covers, tv-jingles, weird fringe stuff.</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>
        """)
    
    initial_load_done = gr.State(value=False)
    uploaded_file_state = gr.State()
    
    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 PREVIEW ==========
        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/Lyrics",
                        value="",
                        lines=5
                    )
                    
                    cover_title = gr.Textbox(
                        label="Title",
                        value=""
                    )
                    
                    cover_style = gr.Textbox(
                        label="Style",
                        value=""
                    )
                    
                    gr.Markdown("### 🔗 Audio Source")
                    
                    # Radio button to choose source
                    use_custom_url = gr.Radio(
                        label="Select source type",
                        choices=[
                            ("Use uploaded file", False),
                            ("Use custom URL", True)
                        ],
                        value=False,
                        info="Choose where your audio comes from"
                    )
                    
                    # Custom URL input (visible when custom is selected)
                    custom_url_input = gr.Textbox(
                        label="Custom URL",
                        placeholder="https://example.com/audio.mp3 or /file=/path/to/file",
                        visible=False
                    )
                    
                    # Preview section
                    gr.Markdown("### 🎧 Audio Preview")
                    gr.Markdown("Listen to verify your audio source works:")
                    
                    preview_audio = gr.Audio(
                        label="Preview",
                        type="filepath",
                        interactive=False,
                        visible=False
                    )
                    
                    # Store uploaded file reference
                    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"
                    )
            
            # Update visibility of custom URL input
            def toggle_custom_url(choice):
                return gr.update(visible=choice)
            
            use_custom_url.change(
                toggle_custom_url,
                inputs=use_custom_url,
                outputs=custom_url_input
            )
            
            # Update preview when source changes
            def update_preview_with_source(use_custom, custom_url, uploaded):
                if use_custom and custom_url.strip():
                    return gr.update(value=custom_url.strip(), visible=True)
                elif not use_custom and uploaded:
                    file_path = get_file_path(uploaded)
                    if file_path:
                        return gr.update(value=file_path, visible=True)
                return gr.update(value=None, visible=False)
            
            # Trigger preview updates
            use_custom_url.change(
                fn=update_preview_with_source,
                inputs=[use_custom_url, custom_url_input, uploaded_file_ref],
                outputs=preview_audio
            )
            
            custom_url_input.change(
                fn=update_preview_with_source,
                inputs=[use_custom_url, custom_url_input, uploaded_file_ref],
                outputs=preview_audio
            )
            
            # 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 Any Task", id="tab_check"):
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("### Check Task Status")
                    
                    check_task_id = gr.Textbox(
                        label="Task ID",
                        placeholder="Enter Task ID from generation"
                    )
                    
                    check_btn = gr.Button("🔍 Check Status", variant="primary")
                    check_clear_btn = gr.Button("🗑️ Clear", variant="secondary")
                    
                    gr.Markdown("""
                    **Quick access via URL:**
                    Add `?taskid=YOUR_TASK_ID` to the URL
                    
                    Example:
                    `https://1hit.no/gen/view.php?task_id=fa3529d5cbaa93427ee4451976ed5c4b`
                    """)
                
                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."
                    )
        
        # ========== PAGE 4: CALLBACK LOGS ==========
        with gr.TabItem("📋 4. Callback Logs", id="tab_logs"):
            gr.Markdown(f"""
            ### 📋 Callback Logs
            View all callback data from Suno API at: [`{LOGS_URL}`]({LOGS_URL})
            """)
            
            logs_iframe = gr.HTML(f"""
            <iframe src="{LOGS_URL}" class="logs-iframe" allow="accelerometer; autoplay; clipboard-write; encrypted-media; gyroscope; picture-in-picture" allowfullscreen></iframe>
            """)
            
            with gr.Row():
                refresh_logs_btn = gr.Button("🔄 Refresh Logs", size="sm")
                
                def refresh_logs():
                    return f"""
                    <iframe src="{LOGS_URL}?t={int(time.time())}" class="logs-iframe" allow="accelerometer; autoplay; clipboard-write; encrypted-media; gyroscope; picture-in-picture" allowfullscreen></iframe>
                    """
                
                refresh_logs_btn.click(
                    fn=refresh_logs,
                    outputs=logs_iframe
                )

        # ========== PAGE 4: CALLBACK LOGS ==========
        with gr.TabItem("📋 5. Backup/Download", id="tab_backup"):
            gr.Markdown(f"""
            ### 📋 Callback Logs
            View all callback data from Suno API at: [`{LOGS_URL}`]({LOGS_URL})
            """)
            
            logs_iframe = gr.HTML(f"""
            <iframe src="{BACK_URL}" class="logs-iframe" allow="accelerometer; autoplay; clipboard-write; encrypted-media; gyroscope; picture-in-picture" allowfullscreen></iframe>
            """)
            
            with gr.Row():
                refresh_logs_btn = gr.Button("🔄 Refresh Logs", size="sm")
                
                def refresh_logs():
                    return f"""
                    <iframe src="{BACK_URL}?t={int(time.time())}" class="logs-iframe" allow="accelerometer; autoplay; clipboard-write; encrypted-media; gyroscope; picture-in-picture" allowfullscreen></iframe>
                    """
                
                refresh_logs_btn.click(
                    fn=refresh_logs,
                    outputs=logs_iframe
                )
    
    # ============================================
    # EVENT HANDLERS
    # ============================================
    
    # Page 1 handlers
    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]
    )
    
    def store_file_ref(file_obj):
        return file_obj
    
    file_input.change(
        fn=store_file_ref,
        inputs=[file_input],
        outputs=uploaded_file_ref
    )
    
    # Also update preview when new file is uploaded
    file_input.change(
        fn=lambda f, use_custom, custom: update_preview_with_source(use_custom, custom, f),
        inputs=[file_input, use_custom_url, custom_url_input],
        outputs=preview_audio
    )
    
    # Store uploaded file for Page 2
    file_input.change(
        fn=store_file_ref,
        inputs=[file_input],
        outputs=uploaded_file_state
    )
    
    # Page 2 generate handler
    generate_btn.click(
        fn=generate_cover,
        inputs=[
            cover_prompt, cover_title, cover_style,
            use_custom_url, custom_url_input, uploaded_file_state,
            cover_instrumental, cover_model,
            gr.State(""), cover_negative, cover_vocal,
            gr.State(0.65), gr.State(0.65), gr.State(0.65),
            gr.State(True)
        ],
        outputs=generate_output
    )
    
    # Page 3 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
    )
    
    # URL parameter handling
    app.load(
        fn=on_page_load,
        inputs=[],
        outputs=[check_task_id, check_output, tabs, initial_load_done],
        queue=False
    )
    
    # Footer
    gr.Markdown("---")
    gr.Markdown(f"""
    <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>
        <p>📋 <b>Logs:</b> <a href="{LOGS_URL}" target="_blank">{LOGS_URL}</a></p>
        <p><b>💡 Preview Feature:</b> Listen to your audio before generating to verify the URL works!</p>
    </div>
    """)

if __name__ == "__main__":
    print(f"🚀 Starting Suno Cover Creator with Audio Preview")
    print(f"🔑 API Key: {'✅ Set' if SUNO_API_KEY else '❌ Not set'}")
    print(f"📋 Logs: {LOGS_URL}")
    
    app.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )