import gradio as gr
import requests
import os
import json
from datetime import datetime
from typing import List, Tuple

# Configuration
OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY", "")
OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1"

# Available models with details
MODELS = {
    "OpenAI": {
        "gpt-4-turbo-preview": "Latest GPT-4 Turbo",
        "gpt-4": "GPT-4 Standard", 
        "gpt-3.5-turbo": "GPT-3.5 Turbo (Fastest)",
        "gpt-3.5-turbo-16k": "GPT-3.5 Turbo 16K"
    },
    "Anthropic": {
        "anthropic/claude-3-opus-20240229": "Claude 3 Opus (Most Powerful)",
        "anthropic/claude-3-sonnet-20240229": "Claude 3 Sonnet (Balanced)",
        "anthropic/claude-3-haiku-20240307": "Claude 3 Haiku (Fastest)",
        "anthropic/claude-2.1": "Claude 2.1",
        "anthropic/claude-instant-1.2": "Claude Instant"
    },
    "Google": {
        "google/gemini-pro": "Gemini Pro",
        "google/palm-2-chat-bison": "PaLM 2 Chat"
    },
    "Meta": {
        "meta-llama/llama-3-70b-instruct": "Llama 3 70B",
        "meta-llama/llama-3-8b-instruct": "Llama 3 8B",
        "meta-llama/llama-2-70b-chat": "Llama 2 70B"
    },
    "Mistral": {
        "mistralai/mistral-7b-instruct": "Mistral 7B",
        "mistralai/mixtral-8x7b-instruct": "Mixtral 8x7B"
    },
    "Others": {
        "databricks/dbrx-instruct": "DBRX Instruct",
        "cohere/command-r-plus": "Command R+",
        "zero-one-ai/yi-34b-chat": "Yi 34B Chat"
    }
}

# Flatten models for dropdown
MODEL_OPTIONS = []
for provider, models in MODELS.items():
    for model_id, model_name in models.items():
        MODEL_OPTIONS.append(f"{provider}: {model_name}")

# Custom CSS
custom_css = """
.gradio-container {
    font-family: 'Inter', sans-serif;
}

.main-container {
    max-width: 1200px;
    margin: 0 auto;
}

/* Chat container styling */
.chat-container {
    height: 600px !important;
    overflow-y: auto;
    border-radius: 10px;
    padding: 20px;
    background-color: var(--background-fill-secondary);
}

/* Message styling */
.message {
    margin: 10px 0;
    padding: 15px;
    border-radius: 10px;
    animation: fadeIn 0.3s ease-in;
}

.user-message {
    background-color: var(--primary-500);
    color: white;
    margin-left: 20%;
}

.assistant-message {
    background-color: var(--neutral-100);
    margin-right: 20%;
}

/* Input area styling */
.input-container {
    border-top: 2px solid var(--border-color-primary);
    padding-top: 20px;
    margin-top: 20px;
}

/* Model selector styling */
.model-selector {
    background-color: var(--background-fill-secondary);
    padding: 10px;
    border-radius: 8px;
    margin-bottom: 10px;
}

/* Animation */
@keyframes fadeIn {
    from { opacity: 0; transform: translateY(10px); }
    to { opacity: 1; transform: translateY(0); }
}

/* Status indicator */
.status-indicator {
    display: inline-block;
    width: 10px;
    height: 10px;
    border-radius: 50%;
    margin-right: 5px;
}

.status-online { background-color: #10b981; }
.status-offline { background-color: #ef4444; }
.status-loading { background-color: #f59e0b; }
"""

def get_model_id_from_name(model_name: str) -> str:
    """Extract model ID from display name"""
    for provider, models in MODELS.items():
        for model_id, name in models.items():
            if f"{provider}: {name}" == model_name:
                return model_id
    return "gpt-3.5-turbo"  # fallback

def check_api_status() -> bool:
    """Check if API key is set and valid"""
    return bool(OPENROUTER_API_KEY)

def format_message_history(history: List[Tuple[str, str]]) -> List[dict]:
    """Format chat history for API request"""
    messages = []
    for user_msg, assistant_msg in history:
        if user_msg:
            messages.append({"role": "user", "content": user_msg})
        if assistant_msg:
            messages.append({"role": "assistant", "content": assistant_msg})
    return messages

def call_openrouter_api(
    message: str, 
    history: List[Tuple[str, str]], 
    model_name: str, 
    temperature: float,
    max_tokens: int,
    system_prompt: str
) -> str:
    """Call OpenRouter API"""
    
    if not check_api_status():
        return "❌ API Key not set. Please set OPENROUTER_API_KEY in Space Settings > Variables and Secrets"
    
    model_id = get_model_id_from_name(model_name)
    
    headers = {
        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
        "Content-Type": "application/json",
        "HTTP-Referer": "https://huggingface.co",
        "X-Title": "HuggingFace Space"
    }
    
    # Build messages
    messages = []
    if system_prompt:
        messages.append({"role": "system", "content": system_prompt})
    
    messages.extend(format_message_history(history))
    messages.append({"role": "user", "content": message})
    
    data = {
        "model": model_id,
        "messages": messages,
        "temperature": temperature,
        "max_tokens": max_tokens,
        "stream": True
    }
    
    try:
        response = requests.post(
            f"{OPENROUTER_BASE_URL}/chat/completions",
            headers=headers,
            json=data,
            stream=True
        )
        
        if response.status_code != 200:
            return f"❌ Error: {response.status_code} - {response.text}"
        
        # Handle streaming response
        full_response = ""
        for line in response.iter_lines():
            if line:
                line_str = line.decode('utf-8')
                if line_str.startswith("data: "):
                    if line_str == "data: [DONE]":
                        break
                    try:
                        chunk = json.loads(line_str[6:])
                        if 'choices' in chunk and chunk['choices']:
                            delta = chunk['choices'][0].get('delta', {})
                            content = delta.get('content', '')
                            full_response += content
                            yield full_response
                    except json.JSONDecodeError:
                        continue
                        
    except requests.exceptions.RequestException as e:
        yield f"❌ Connection error: {str(e)}"
    except Exception as e:
        yield f"❌ Unexpected error: {str(e)}"

# Create Gradio interface
with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
    # Header
    gr.Markdown(
        """
        # 🌟 OpenRouter Multi-Model Chat
        
        <div style='text-align: center; margin-bottom: 20px;'>
            <p>Chat with various AI models including GPT-4, Claude 3, Gemini, and more!</p>
            <p>Powered by <a href='https://openrouter.ai' target='_blank'>OpenRouter</a></p>
        </div>
        """
    )
    
    # API Status
    with gr.Row():
        with gr.Column(scale=1):
            api_status = gr.Markdown(
                f"**API Status:** {'🟢 Connected' if check_api_status() else '🔴 Not Connected'}"
            )
    
    # Main chat interface
    with gr.Row():
        # Left sidebar - Settings
        with gr.Column(scale=3):
            model_dropdown = gr.Dropdown(
                choices=MODEL_OPTIONS,
                value="OpenAI: GPT-3.5 Turbo (Fastest)",
                label="🤖 Select Model",
                interactive=True
            )
            
            with gr.Accordion("⚙️ Advanced Settings", open=False):
                system_prompt = gr.Textbox(
                    label="System Prompt",
                    placeholder="You are a helpful assistant...",
                    lines=3
                )
                
                temperature = gr.Slider(
                    minimum=0,
                    maximum=2,
                    value=0.7,
                    step=0.1,
                    label="Temperature (Creativity)"
                )
                
                max_tokens = gr.Slider(
                    minimum=50,
                    maximum=4096,
                    value=1024,
                    step=50,
                    label="Max Tokens"
                )
            
            # Model info
            gr.Markdown(
                """
                ### 📊 Model Categories
                - **OpenAI**: GPT-4, GPT-3.5
                - **Anthropic**: Claude 3 (Opus, Sonnet, Haiku)
                - **Google**: Gemini Pro, PaLM 2
                - **Meta**: Llama 3, Llama 2
                - **Mistral**: Mistral 7B, Mixtral
                """
            )
        
        # Right side - Chat
        with gr.Column(scale=7):
            chatbot = gr.Chatbot(
                height=500,
                label="💬 Conversation",
                elem_classes="chat-container"
            )
            
            with gr.Row():
                msg = gr.Textbox(
                    label="Type your message",
                    placeholder="Ask me anything...",
                    lines=2,
                    scale=9
                )
                
                with gr.Column(scale=1, min_width=80):
                    send_btn = gr.Button("Send 📤", variant="primary")
                    clear_btn = gr.Button("Clear 🗑️")
    
    # Examples
    gr.Examples(
        examples=[
            "What are the key differences between GPT-4 and Claude 3?",
            "Write a Python function to calculate fibonacci numbers",
            "Explain quantum computing in simple terms",
            "What's the meaning of life?",
            "Help me write a professional email"
        ],
        inputs=msg,
        label="💡 Example Prompts"
    )
    
    # Event handlers
    def respond(message, chat_history, model, temp, tokens, sys_prompt):
        bot_message = ""
        for response in call_openrouter_api(
            message, chat_history, model, temp, tokens, sys_prompt
        ):
            bot_message = response
            yield chat_history + [(message, bot_message)]
    
    def clear_chat():
        return None
    
    # Connect events
    msg.submit(
        respond,
        [msg, chatbot, model_dropdown, temperature, max_tokens, system_prompt],
        chatbot
    ).then(
        lambda: "", None, msg
    )
    
    send_btn.click(
        respond,
        [msg, chatbot, model_dropdown, temperature, max_tokens, system_prompt],
        chatbot
    ).then(
        lambda: "", None, msg
    )
    
    clear_btn.click(clear_chat, None, chatbot)
    
    # Footer
    gr.Markdown(
        """
        ---
        <div style='text-align: center; color: gray;'>
            Made with ❤️ using Gradio | 
            <a href='https://github.com/yourusername/yourrepo' target='_blank'>GitHub</a> | 
            <a href='https://openrouter.ai/docs' target='_blank'>API Docs</a>
        </div>
        """
    )

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