import os
import requests
import gradio as gr
from datetime import datetime

# API Configuration
API_KEY = os.getenv('StableCogKey')
API_HOST = 'https://api.stablecog.com'

headers = {
    'Authorization': f'Bearer {API_KEY}',
    'Content-Type': 'application/json'
}

# Global state for pagination
current_outputs = []
current_page = 0
page_size = 10

def get_models():
    """Fetch and display available models"""
    try:
        url = f'{API_HOST}/v1/image/generation/models'
        response = requests.get(url, headers=headers, timeout=10)
        
        if response.status_code == 200:
            data = response.json()
            models = data.get('models', [])
            
            # Format display
            display_text = f"📊 Found {len(models)} models\n"
            display_text += f"⏰ {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n"
            
            for i, model in enumerate(models, 1):
                name = model.get('name', 'Unknown')
                model_type = model.get('type', 'unknown')
                description = model.get('description', 'No description')
                
                display_text += f"{i}. 🔹 **{name}**\n"
                display_text += f"   📝 {description}\n"
                display_text += f"   🏷️ Type: {model_type}\n"
                display_text += f"   🌍 Public: {'✅' if model.get('is_public') else '❌'}\n"
                display_text += f"   ⭐ Default: {'✅' if model.get('is_default') else '❌'}\n"
                display_text += f"   👥 Community: {'✅' if model.get('is_community') else '❌'}\n"
                display_text += "─" * 40 + "\n"
            
            return display_text, str(data)
            
        else:
            return f"❌ Error {response.status_code}", f"Error: {response.text}"
            
    except Exception as e:
        return f"❌ Error: {str(e)}", "No data"

def fetch_outputs():
    """Fetch outputs from API"""
    global current_outputs
    try:
        url = f'{API_HOST}/v1/image/generation/outputs'
        response = requests.get(url, headers=headers, timeout=10)
        
        if response.status_code == 200:
            data = response.json()
            current_outputs = data.get('outputs', [])
            return data
        else:
            return None
    except:
        return None

def update_outputs_display():
    """Update outputs display with current page"""
    global current_outputs, current_page, page_size
    
    if not current_outputs:
        return "📭 No outputs found. Generate some images first!", "<div style='text-align: center; padding: 40px; color: #888;'>No images found</div>", "[]"
    
    total = len(current_outputs)
    total_pages = (total + page_size - 1) // page_size  # Ceiling division
    
    # Calculate page bounds
    start_idx = current_page * page_size
    end_idx = min(start_idx + page_size, total)
    page_outputs = current_outputs[start_idx:end_idx]
    
    # Format display
    display_text = f"📊 Total outputs: {total}\n"
    display_text += f"📄 Page {current_page + 1} of {total_pages}\n"
    display_text += f"🖼️ Showing {start_idx + 1}-{end_idx} of {total}\n"
    display_text += f"⏰ {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n"
    
    # Create gallery HTML with lightbox
    gallery_html = """
    <style>
        .image-gallery {
            display: grid;
            grid-template-columns: repeat(auto-fill, minmax(200px, 1fr));
            gap: 15px;
            margin-bottom: 20px;
        }
        .image-card {
            border-radius: 10px;
            overflow: hidden;
            background: rgba(255,255,255,0.1);
            padding: 10px;
            cursor: pointer;
            transition: transform 0.2s;
        }
        .image-card:hover {
            transform: scale(1.02);
            background: rgba(255,255,255,0.15);
        }
        .image-card img {
            width: 100%;
            height: 200px;
            object-fit: cover;
            border-radius: 8px;
        }
        .image-meta {
            margin-top: 8px;
            font-size: 12px;
            line-height: 1.4;
        }
        .lightbox {
            display: none;
            position: fixed;
            z-index: 9999;
            left: 0;
            top: 0;
            width: 100%;
            height: 100%;
            background: rgba(0,0,0,0.9);
            justify-content: center;
            align-items: center;
        }
        .lightbox img {
            max-width: 90%;
            max-height: 90%;
            border-radius: 10px;
            box-shadow: 0 20px 60px rgba(0,0,0,0.5);
        }
        .lightbox.active {
            display: flex;
        }
        .close-btn {
            position: absolute;
            top: 20px;
            right: 30px;
            color: white;
            font-size: 40px;
            font-weight: bold;
            cursor: pointer;
            z-index: 10000;
        }
        .pagination {
            display: flex;
            justify-content: center;
            gap: 10px;
            margin-top: 20px;
        }
        .page-btn {
            padding: 8px 16px;
            background: rgba(255,255,255,0.1);
            border: 1px solid rgba(255,255,255,0.2);
            border-radius: 8px;
            color: white;
            cursor: pointer;
            transition: background 0.2s;
        }
        .page-btn:hover {
            background: rgba(255,255,255,0.2);
        }
        .page-btn.disabled {
            opacity: 0.5;
            cursor: not-allowed;
        }
        .page-info {
            display: flex;
            align-items: center;
            padding: 8px 16px;
            color: white;
        }
    </style>
    
    <div class="lightbox" id="lightbox" onclick="closeLightbox()">
        <span class="close-btn" onclick="closeLightbox()">&times;</span>
        <img id="lightbox-img" onclick="event.stopPropagation()">
    </div>
    
    <script>
        function openLightbox(imgSrc) {
            document.getElementById('lightbox-img').src = imgSrc;
            document.getElementById('lightbox').classList.add('active');
        }
        function closeLightbox() {
            document.getElementById('lightbox').classList.remove('active');
        }
        // Close lightbox on ESC key
        document.addEventListener('keydown', function(e) {
            if (e.key === 'Escape') closeLightbox();
        });
    </script>
    
    <div class="image-gallery">
    """
    
    for idx, output in enumerate(page_outputs, start=start_idx + 1):
        output_id = output.get('id', 'N/A')
        created_at = output.get('created_at', 'N/A')
        model_name = output.get('model_name', 'Unknown')
        
        # Gallery status with emojis
        gallery_status = output.get('gallery_status', 'not_submitted')
        gallery_emoji = {
            'not_submitted': '🔒',
            'submitted': '📤',
            'approved': '✅',
            'rejected': '❌'
        }.get(gallery_status, '❓')
        
        # Favorites
        is_favorited = output.get('is_favorited', False)
        favorite_emoji = '❤️' if is_favorited else '🤍'
        
        # Format timestamp
        if created_at != 'N/A':
            try:
                dt = datetime.fromisoformat(created_at.replace('Z', '+00:00'))
                created_date = dt.strftime('%Y-%m-%d')
            except:
                created_date = created_at
        else:
            created_date = 'Unknown'
        
        # Get image URL
        image_url = output.get('image_url')
        if not image_url:
            image_urls = output.get('image_urls', [])
            if image_urls and isinstance(image_urls, list) and len(image_urls) > 0:
                image_url = image_urls[0]
        
        if image_url:
            # Text display
            display_text += f"{idx}. 🔹 **Output {output_id[:8]}...**\n"
            display_text += f"   🕒 {created_at}\n"
            display_text += f"   🤖 {model_name}\n"
            display_text += f"   🖼️ {gallery_emoji} {gallery_status}\n"
            display_text += f"   {favorite_emoji} Favorite\n"
            
            # Show generation details if available
            generation = output.get('generation', {})
            if generation:
                prompt = generation.get('prompt', 'No prompt')
                if len(prompt) > 50:
                    prompt = prompt[:50] + '...'
                display_text += f"   📝 {prompt}\n"
            
            display_text += "─" * 40 + "\n"
            
            # Add to gallery with click event
            gallery_html += f"""
            <div class="image-card" onclick="openLightbox('{image_url}')">
                <img src="{image_url}" alt="Output {idx}">
                <div class="image-meta">
                    <div>#{idx} 📅 {created_date}</div>
                    <div>🤖 {model_name[:15]}{'...' if len(model_name) > 15 else ''}</div>
                    <div>{gallery_emoji} {gallery_status}</div>
                    <div>{favorite_emoji}</div>
                </div>
            </div>
            """
        else:
            display_text += f"{idx}. ⚠️ No image data\n"
            display_text += "─" * 40 + "\n"
    
    gallery_html += "</div>"
    
    # Add pagination controls
    gallery_html += f"""
    <div class="pagination">
        <button class="page-btn {'disabled' if current_page == 0 else ''}" 
                onclick="{'return false;' if current_page == 0 else f'window.paginationButtonClick({current_page - 1})'}">
            ◀ Previous
        </button>
        <div class="page-info">
            Page {current_page + 1} of {total_pages}
        </div>
        <button class="page-btn {'disabled' if current_page >= total_pages - 1 else ''}" 
                onclick="{'return false;' if current_page >= total_pages - 1 else f'window.paginationButtonClick({current_page + 1})'}">
            Next ▶
        </button>
    </div>
    """
    
    return display_text, gallery_html, str(current_outputs)

def load_outputs():
    """Load outputs from API and display first page"""
    global current_page
    current_page = 0
    data = fetch_outputs()
    if data:
        return update_outputs_display()
    else:
        return "❌ Failed to load outputs", "<div style='color: red; padding: 20px;'>Failed to load outputs</div>", "[]"

def next_page():
    """Go to next page"""
    global current_page
    if current_outputs:
        total_pages = (len(current_outputs) + page_size - 1) // page_size
        if current_page < total_pages - 1:
            current_page += 1
    return update_outputs_display()

def prev_page():
    """Go to previous page"""
    global current_page
    if current_page > 0:
        current_page -= 1
    return update_outputs_display()

# Create Gradio interface
with gr.Blocks(title="StableCog Dashboard") as demo:
    gr.Markdown("# 🎯 StableCog Dashboard")
    
    with gr.Tabs():
        with gr.Tab("🤖 Models"):
            with gr.Row():
                models_display = gr.Textbox(label="Available Models", lines=25)
                models_raw = gr.Code(label="Raw JSON", language="json", lines=25)
            
            check_models_btn = gr.Button("🔄 Refresh Models", variant="primary")
            check_models_btn.click(get_models, outputs=[models_display, models_raw])
        
        with gr.Tab("🖼️ Outputs"):
            with gr.Row():
                with gr.Column(scale=1):
                    outputs_display = gr.Textbox(label="Output Details", lines=25)
                with gr.Column(scale=2):
                    outputs_gallery = gr.HTML(label="Image Gallery")
            
            with gr.Row():
                outputs_raw = gr.Code(label="Raw JSON", language="json", lines=15)
            
            with gr.Row():
                load_outputs_btn = gr.Button("🔄 Load Outputs", variant="primary")
                prev_page_btn = gr.Button("◀ Previous Page")
                next_page_btn = gr.Button("Next Page ▶")
            
            # JavaScript for pagination
            js = """
            <script>
                window.paginationButtonClick = function(page) {
                    const event = new CustomEvent('gradio_pagination', { detail: { page: page } });
                    document.dispatchEvent(event);
                }
            </script>
            """
            gr.HTML(js)
            
            # Connect buttons
            load_outputs_btn.click(
                load_outputs, 
                outputs=[outputs_display, outputs_gallery, outputs_raw]
            )
            
            prev_page_btn.click(
                prev_page,
                outputs=[outputs_display, outputs_gallery, outputs_raw]
            )
            
            next_page_btn.click(
                next_page,
                outputs=[outputs_display, outputs_gallery, outputs_raw]
            )

if __name__ == "__main__":
    demo.launch()