import os
import requests
import gradio as gr

# API Configuration
API_KEY = os.getenv('StableCogKey')
API_HOST = 'https://api.stablecog.com'
MODELS_ENDPOINT = '/v1/image/generation/models'
URL = f'{API_HOST}{MODELS_ENDPOINT}'

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

def get_models():
    """Fetch and display available models"""
    try:
        response = requests.get(URL, headers=headers, timeout=10)
        
        if response.status_code == 200:
            data = response.json()
            models = data.get('models', [])
            
            # Format display
            display_text = ""
            for model in models:
                name = model.get('name', 'Unknown')
                model_type = model.get('type', 'unknown')
                description = model.get('description', 'No description')
                
                display_text += f"🔹 {name} ({model_type})\n"
                display_text += f"   📝 {description}\n"
                display_text += f"   📊 Public: {model.get('is_public', False)}\n"
                display_text += f"   ⭐ Default: {model.get('is_default', False)}\n"
                display_text += f"   👥 Community: {model.get('is_community', False)}\n"
                display_text += "─" * 40 + "\n"
            
            total = len(models)
            display_text = f"📊 Found {total} models:\n\n" + display_text
            
            return display_text, str(data)
            
        else:
            return f"❌ Error: {response.status_code}", f"Error: {response.text}"
            
    except Exception as e:
        return f"❌ Connection error: {str(e)}", "No data"

# Create Gradio interface
with gr.Blocks(title="StableCog Models") as demo:
    gr.Markdown("# 🤖 StableCog Models Checker")
    
    with gr.Row():
        models_display = gr.Textbox(label="Models", lines=20)
        raw_output = gr.Code(label="Raw JSON", language="json", lines=20)
    
    check_btn = gr.Button("🔄 Check Models", variant="primary")
    check_btn.click(get_models, outputs=[models_display, raw_output])

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