import streamlit as st
import json
import requests
import time
import pandas as pd
from datetime import datetime
from urllib.parse import urlparse
import os

# Polling function
def poll_status(task_id, single_poll=False):
    """Poll Suno API for task status including task ID"""
    if not task_id:
        st.error("❌ Task ID cannot be empty")
        return False
    
    headers = {
        "Authorization": f"Bearer {st.session_state.suno_api_key}",
        "Content-Type": "application/json"
    }
    
    try:
        with st.spinner(f"Checking status for task: {task_id}..."):
            # Try different endpoint formats
            endpoints_to_try = [
                f"https://api.sunoapi.org/api/v1/wav/record-info?taskId={task_id}",
                f"https://api.sunoapi.org/api/v1/wav/record-info/{task_id}",
                f"https://api.sunoapi.org/api/v1/wav/record-info?id={task_id}"
            ]
            
            response = None
            for endpoint in endpoints_to_try:
                try:
                    response = requests.get(
                        endpoint,
                        headers=headers,
                        timeout=30
                    )
                    if response.status_code == 200:
                        break
                except:
                    continue
            
            if not response:
                st.error("❌ Could not connect to any polling endpoint")
                return False
                
            if response.status_code == 200:
                result = response.json()
                
                # Store in session state
                if 'poll_results' not in st.session_state:
                    st.session_state.poll_results = []
                
                poll_entry = {
                    "timestamp": datetime.now().strftime("%H:%M:%S"),
                    "task_id": task_id,
                    "response": result
                }
                st.session_state.poll_results.append(poll_entry)
                
                # Display result
                if result.get("code") == 200:
                    data = result.get("data", {})
                    
                    if data.get("successFlag") == "SUCCESS":
                        # Audio is ready!
                        audio_url = data.get("response", {}).get("audioWavUrl")
                        if audio_url:
                            st.success("✅ Audio ready!")
                            st.balloons()
                            
                            # Display audio info
                            st.markdown("### 🎵 Audio Information")
                            cols = st.columns(2)
                            with cols[0]:
                                st.metric("Task ID", data.get("taskId", "N/A")[:12] + "...")
                                st.metric("Music ID", data.get("musicId", "N/A")[:12] + "...")
                            with cols[1]:
                                st.metric("Status", data.get("successFlag", "N/A"))
                                st.metric("Created", data.get("createTime", "N/A"))
                            
                            # Audio player
                            st.markdown("### 🔊 Listen")
                            st.audio(audio_url, format="audio/wav")
                            
                            # Download section
                            st.markdown("### 📥 Download")
                            st.code(audio_url)
                            
                            col_dl1, col_dl2 = st.columns(2)
                            with col_dl1:
                                st.download_button(
                                    label="⬇ Download WAV",
                                    data=requests.get(audio_url).content if audio_url.startswith("http") else b"",
                                    file_name=f"audio_{data.get('taskId', 'unknown')[:8]}.wav",
                                    mime="audio/wav",
                                    key=f"download_{task_id}"
                                )
                            with col_dl2:
                                if st.button("📋 Copy URL", key=f"copy_{task_id}"):
                                    st.code(audio_url)
                                    st.success("URL copied to clipboard!")
                            
                            # Stop auto-polling if active
                            if st.session_state.get('auto_poll'):
                                st.session_state.auto_poll = False
                                st.success("✅ Auto-polling stopped (audio ready)")
                        else:
                            st.info("⏳ Audio still processing...")
                            st.json(data)
                    else:
                        # Still processing or error
                        status = data.get("successFlag", "UNKNOWN")
                        if status == "PROCESSING":
                            st.info(f"⏳ Processing... ({status})")
                        elif status == "FAILED":
                            st.error(f"❌ Processing failed: {data.get('errorMessage', 'Unknown error')}")
                        else:
                            st.info(f"ℹ️ Status: {status}")
                        
                        # Show progress info
                        st.json(data)
                else:
                    st.error(f"❌ API Error: {result.get('msg', 'Unknown error')}")
                    st.json(result)
            else:
                st.error(f"❌ Status check failed: {response.status_code}")
                try:
                    error_data = response.json()
                    st.json(error_data)
                except:
                    st.text(f"Response: {response.text}")
                    
    except Exception as e:
        st.error(f"⚠️ Polling error: {str(e)}")

def load_callback_logs():
    """Load and parse callback_log.json from remote URL"""
    try:
        # Try to fetch from your callback URL
        log_url = "https://1hit.no/wav/callback_log.json"
        response = requests.get(log_url, timeout=10)
        
        if response.status_code == 200:
            logs = response.json()
            return logs
        else:
            st.warning(f"Could not fetch logs from {log_url}. Status: {response.status_code}")
            return []
    except Exception as e:
        st.warning(f"Could not load callback logs: {str(e)}")
        return []

def parse_callback_logs(logs):
    """Parse callback logs into a structured format"""
    parsed_logs = []
    
    for log in logs:
        # Handle both old and new format
        callback_data = log.get('callback_data') or log.get('data_received') or {}
        
        # Extract data with fallbacks for different formats
        code = callback_data.get('code', 0)
        msg = callback_data.get('msg', '')
        data = callback_data.get('data', {})
        
        # Get audio URL (handle both formats)
        audio_url = data.get('audio_wav_url') or data.get('audioWavUrl') or ''
        task_id = data.get('task_id') or data.get('taskId') or ''
        
        # Get timestamp
        timestamp = log.get('timestamp', 'Unknown')
        
        parsed_logs.append({
            'timestamp': timestamp,
            'code': code,
            'message': msg,
            'task_id': task_id,
            'audio_url': audio_url,
            'raw_data': callback_data
        })
    
    return parsed_logs

# Page configuration
st.set_page_config(
    page_title="Suno API Generator",
    page_icon="🎵",
    layout="wide"
)

# Load YOUR secret names
SUNO_API_KEY = st.secrets.get("SunoKey", "")
CALLBACK_URL = "https://1hit.no/wav/cb.php"

# Initialize session state
if 'task_id' not in st.session_state:
    st.session_state.task_id = None
if 'polling' not in st.session_state:
    st.session_state.polling = False
if 'poll_results' not in st.session_state:
    st.session_state.poll_results = []
if 'auto_poll' not in st.session_state:
    st.session_state.auto_poll = False
if 'suno_api_key' not in st.session_state:
    st.session_state.suno_api_key = SUNO_API_KEY
if 'selected_audio' not in st.session_state:
    st.session_state.selected_audio = None

# Title
st.title("🎵 Suno API Audio Generator")

# Create tabs
tab1, tab2 = st.tabs(["🚀 Generate & Poll", "📊 Callback Logs"])

with tab1:
    # Info section
    with st.expander("📋 How it works", expanded=True):
        st.markdown("""
        ### Workflow:
        1. **Step 1**: Submit task to Suno API → Get `taskId` from response
        2. **Step 2**: Suno processes audio (async)
        3. **Step 3**: Suno sends callback to: `https://1hit.no/wav/cb.php`
        4. **Step 4**: Check callback logs in the 📊 Callback Logs tab
        5. **Step 5**: OR use Poll button to check status via API
        ### Polling Endpoint:
        ```javascript
        fetch('https://api.sunoapi.org/api/v1/wav/record-info', {
          method: 'GET',
          headers: {
            'Authorization': 'Bearer YOUR_TOKEN'
          }
        })
        ```
        """)

    col1, col2 = st.columns([2, 1])

    with col1:
        st.header("🚀 Step 1: Submit to Suno API")
        
        # Form inputs
        task_id = st.text_input("Task ID", value="5c79****be8e", help="Enter your Suno task ID", key="input_task_id")
        audio_id = st.text_input("Audio ID", value="e231****-****-****-****-****8cadc7dc", help="Enter your Suno audio ID", key="input_audio_id")
        
        if st.button("🎵 Generate Audio", type="primary", use_container_width=True, key="generate_btn"):
            if not SUNO_API_KEY:
                st.error("❌ SunoKey not configured in Hugging Face secrets")
                st.info("Add `SunoKey` secret in Hugging Face Space settings")
                st.stop()
            
            if not task_id or not audio_id:
                st.error("❌ Please enter both Task ID and Audio ID")
                st.stop()
            
            with st.spinner("Sending request to Suno API..."):
                headers = {
                    "Authorization": f"Bearer {SUNO_API_KEY}",
                    "Content-Type": "application/json"
                }
                
                payload = {
                    "taskId": task_id,
                    "audioId": audio_id,
                    "callBackUrl": CALLBACK_URL
                }
                
                try:
                    response = requests.post(
                        "https://api.sunoapi.org/api/v1/wav/generate",
                        headers=headers,
                        json=payload,
                        timeout=30
                    )
                    
                    if response.status_code == 200:
                        result = response.json()
                        
                        # Display response
                        with st.expander("📋 API Response", expanded=True):
                            st.json(result)
                        
                        if result.get("code") == 200 and "data" in result and "taskId" in result["data"]:
                            st.session_state.task_id = result["data"]["taskId"]
                            st.success(f"✅ Task submitted successfully!")
                            st.info(f"**Task ID:** `{st.session_state.task_id}`")
                            
                            # Auto-fill poll input
                            st.session_state.poll_task_input = st.session_state.task_id
                            
                            st.markdown("---")
                            st.markdown("### 🔄 Next Steps:")
                            st.markdown("""
                            1. **Automatic Callback**: Suno will send result to your callback URL
                            2. **Manual Poll**: Use the Poll button to check status
                            3. **View Logs**: Check the 📊 Callback Logs tab
                            """)
                        else:
                            st.error("❌ Unexpected response format")
                            
                            # Debug info
                            st.markdown("**Debug Response:**")
                            st.code(str(result))
                    else:
                        st.error(f"❌ API Error: {response.status_code}")
                        try:
                            error_data = response.json()
                            st.json(error_data)
                        except:
                            st.text(f"Response: {response.text}")
                            
                except requests.exceptions.Timeout:
                    st.error("⏰ Request timed out after 30 seconds")
                except requests.exceptions.ConnectionError:
                    st.error("🔌 Connection error - check your internet connection")
                except requests.exceptions.RequestException as e:
                    st.error(f"⚠️ Request failed: {str(e)}")

    with col2:
        st.header("🔍 Step 2: Poll Status")
        
        # Manual task ID input for polling
        poll_task_id = st.text_input(
            "Task ID to Poll",
            value=st.session_state.get('poll_task_input', st.session_state.task_id or ""),
            help="Enter task ID to check status",
            key="poll_task_input_field"
        )
        
        col2a, col2b = st.columns(2)
        
        with col2a:
            poll_once_btn = st.button("🔄 Poll Once", type="secondary", use_container_width=True, key="poll_once_btn")
        
        with col2b:
            auto_poll_btn = st.button("🔁 Auto Poll (10s)", type="secondary", use_container_width=True, key="auto_poll_btn")

    # Handle poll button clicks
    if poll_once_btn:
        if not poll_task_id:
            st.error("❌ Please enter a Task ID")
        elif not SUNO_API_KEY:
            st.error("❌ SunoKey not configured")
        else:
            poll_status(poll_task_id, single_poll=True)

    if auto_poll_btn:
        if not poll_task_id:
            st.error("❌ Please enter a Task ID")
        elif not SUNO_API_KEY:
            st.error("❌ SunoKey not configured")
        else:
            st.session_state.auto_poll = True
            st.session_state.auto_poll_task_id = poll_task_id
            st.rerun()

    # Display poll results if any
    if st.session_state.poll_results:
        st.markdown("---")
        st.header("📊 Polling History")
        
        # Show last 5 results
        recent_results = list(reversed(st.session_state.poll_results[-5:]))
        
        for i, result in enumerate(recent_results):
            with st.expander(f"Poll {i+1} - {result['timestamp']} - Task: {result['task_id'][:12]}...", expanded=(i == len(recent_results)-1)):
                response_data = result["response"]
                
                if response_data.get("code") == 200:
                    data = response_data.get("data", {})
                    status = data.get("successFlag", "UNKNOWN")
                    
                    # Status badge
                    if status == "SUCCESS":
                        st.success(f"✅ {status}")
                    elif status == "PROCESSING":
                        st.info(f"⏳ {status}")
                    elif status == "FAILED":
                        st.error(f"❌ {status}")
                    else:
                        st.warning(f"⚠️ {status}")
                    
                    # Display key info
                    cols = st.columns(3)
                    with cols[0]:
                        st.metric("Task ID", data.get("taskId", "N/A")[:12] + "...")
                    with cols[1]:
                        st.metric("Music ID", data.get("musicId", "N/A")[:12] + "...")
                    with cols[2]:
                        st.metric("Status", status)
                    
                    # Show audio if available
                    audio_url = data.get("response", {}).get("audioWavUrl")
                    if audio_url:
                        st.audio(audio_url, format="audio/wav")
                        st.markdown(f"**Audio URL:** `{audio_url}`")
                    
                    # Show full response
                    with st.expander("📋 Full Response JSON"):
                        st.json(response_data)
                else:
                    st.error(f"❌ Error: {response_data.get('msg', 'Unknown error')}")
                    st.json(response_data)

with tab2:
    st.header("📊 Callback Log Viewer")
    st.markdown("View and manage audio files from callback logs")
    
    # Load callback logs
    with st.spinner("Loading callback logs..."):
        logs = load_callback_logs()
        parsed_logs = parse_callback_logs(logs)
    
    if not parsed_logs:
        st.info("📭 No callback logs found. Callbacks will appear here when Suno API sends them.")
        st.markdown(f"**Callback URL:** `{CALLBACK_URL}`")
    else:
        # Display statistics
        total_logs = len(parsed_logs)
        successful_logs = len([log for log in parsed_logs if log['code'] == 200])
        audio_logs = len([log for log in parsed_logs if log['audio_url']])
        
        col_stats1, col_stats2, col_stats3 = st.columns(3)
        with col_stats1:
            st.metric("📊 Total Callbacks", total_logs)
        with col_stats2:
            st.metric("✅ Successful", successful_logs)
        with col_stats3:
            st.metric("🎵 Audio Files", audio_logs)
        
        # Create a DataFrame for display
        df_data = []
        for log in parsed_logs:
            df_data.append({
                'Timestamp': log['timestamp'],
                'Status': '✅ Success' if log['code'] == 200 else '❌ Error',
                'Task ID': log['task_id'][:12] + '...' if log['task_id'] else 'N/A',
                'Audio URL': log['audio_url'][:50] + '...' if log['audio_url'] else 'No audio',
                'Message': log['message'][:50] + '...' if log['message'] else '',
                'Full Data': log
            })
        
        if df_data:
            df = pd.DataFrame(df_data)
            
            # Search and filter
            st.markdown("### 🔍 Search & Filter")
            search_term = st.text_input("Search by Task ID or Message:", placeholder="Enter search term...")
            
            if search_term:
                filtered_df = df[df.apply(lambda row: search_term.lower() in str(row['Task ID']).lower() or 
                                                    search_term.lower() in str(row['Message']).lower(), axis=1)]
            else:
                filtered_df = df
            
            # Display table
            st.markdown("### 📋 Callback Logs")
            for idx, row in filtered_df.iterrows():
                with st.expander(f"{row['Timestamp']} - {row['Status']} - Task: {row['Task ID']}", expanded=False):
                    col1, col2 = st.columns(2)
                    
                    with col1:
                        st.markdown(f"**Task ID:** `{row['Full Data']['task_id']}`")
                        st.markdown(f"**Status Code:** {row['Full Data']['code']}")
                        st.markdown(f"**Message:** {row['Full Data']['message']}")
                    
                    with col2:
                        if row['Full Data']['audio_url']:
                            st.markdown("**Audio URL:**")
                            st.code(row['Full Data']['audio_url'])
                            
                            # Audio player
                            st.audio(row['Full Data']['audio_url'], format="audio/wav")
                            
                            # Download button
                            audio_filename = row['Full Data']['audio_url'].split('/')[-1] if '/' in row['Full Data']['audio_url'] else f"audio_{row['Full Data']['task_id']}.wav"
                            st.download_button(
                                label="⬇ Download WAV",
                                data=requests.get(row['Full Data']['audio_url']).content if row['Full Data']['audio_url'].startswith("http") else b"",
                                file_name=audio_filename,
                                mime="audio/wav",
                                key=f"download_callback_{idx}"
                            )
                        else:
                            st.info("No audio URL available")
                    
                    # Show full JSON data
                    with st.expander("📋 View Full JSON"):
                        st.json(row['Full Data']['raw_data'])
            
            # Audio player for selected file
            st.markdown("---")
            st.markdown("### 🔊 Audio Player")
            
            # Create a dropdown of available audio files
            audio_files = [log for log in parsed_logs if log['audio_url']]
            
            if audio_files:
                audio_options = [f"{log['timestamp']} - {log['task_id'][:12]}..." for log in audio_files]
                selected_index = st.selectbox(
                    "Select audio to play:",
                    range(len(audio_options)),
                    format_func=lambda x: audio_options[x]
                )
                
                if selected_index is not None:
                    selected_audio = audio_files[selected_index]
                    st.session_state.selected_audio = selected_audio
                    
                    st.markdown(f"**Selected:** {selected_audio['timestamp']} - Task: `{selected_audio['task_id']}`")
                    st.audio(selected_audio['audio_url'], format="audio/wav")
                    
                    # Download button for selected audio
                    audio_filename = selected_audio['audio_url'].split('/')[-1] if '/' in selected_audio['audio_url'] else f"audio_{selected_audio['task_id']}.wav"
                    st.download_button(
                        label=f"⬇ Download {audio_filename}",
                        data=requests.get(selected_audio['audio_url']).content if selected_audio['audio_url'].startswith("http") else b"",
                        file_name=audio_filename,
                        mime="audio/wav",
                        key="download_selected_audio"
                    )
            else:
                st.info("No audio files available in callback logs")
        
        # Refresh button
        if st.button("🔄 Refresh Logs", key="refresh_logs"):
            st.rerun()

# Quick links and info
st.sidebar.markdown("---")
st.sidebar.markdown("### 📝 Quick Links")
st.sidebar.markdown(f"""
- 🔗 [Callback Logs]({CALLBACK_URL.replace('cb.php', 'view.php')})
- 📚 [Suno API Docs](https://docs.sunoapi.org/)
- 🔄 **Poll Endpoint**: `GET /api/v1/wav/record-info`
""")

# Secret status
with st.sidebar.expander("🔐 Secret Status", expanded=False):
    if SUNO_API_KEY:
        masked_key = f"{SUNO_API_KEY[:8]}...{SUNO_API_KEY[-8:]}" if len(SUNO_API_KEY) > 16 else "••••••••"
        st.success(f"✅ SunoKey loaded ({len(SUNO_API_KEY)} chars)")
        st.code(f"SunoKey: {masked_key}")
    else:
        st.error("❌ SunoKey not found")
    
    st.info(f"**Callback URL:** `{CALLBACK_URL}`")

# Auto-polling logic
if st.session_state.get('auto_poll') and st.session_state.get('auto_poll_task_id'):
    task_id = st.session_state.auto_poll_task_id
    
    # Create a placeholder for auto-poll status
    poll_placeholder = st.empty()
    
    with poll_placeholder.container():
        st.info(f"🔁 Auto-polling task: `{task_id[:12]}...` (every 10 seconds)")
        
        # Poll once
        poll_status(task_id, single_poll=True)
        
        # Schedule next poll
        time.sleep(10)
        st.rerun()