import gradio as gr
from huggingface_hub import InferenceClient
import os

MODELS = [
    {"name": "Qwen 2.5 Coder 32B", "id": "Qwen/Qwen2.5-Coder-32B-Instruct"},
    {"name": "Qwen 2.5 72B", "id": "Qwen/Qwen2.5-72B-Instruct"},
    {"name": "Llama 3.3 70B", "id": "meta-llama/Llama-3.3-70B-Instruct"},
    {"name": "Mixtral 8x7B", "id": "mistralai/Mixtral-8x7B-Instruct-v0.1"},
    {"name": "DeepSeek Coder V2", "id": "deepseek-ai/DeepSeek-Coder-V2-Instruct"},
]

def stream_chat(message, history, model_name, api_key):
    if not api_key or api_key.strip() == "":
        yield "⚠️ Please enter your Hugging Face API key in the sidebar first.\n\nGet your free API key at: https://huggingface.co/settings/tokens"
        return
    
    model_id = next((m["id"] for m in MODELS if m["name"] == model_name), MODELS[0]["id"])
    
    try:
        client = InferenceClient(token=api_key.strip())
        
        messages = []
        for h in history:
            if h["role"] == "user":
                messages.append({"role": "user", "content": h["content"]})
            elif h["role"] == "assistant":
                messages.append({"role": "assistant", "content": h["content"]})
        
        messages.append({"role": "user", "content": message})
        
        response = ""
        for chunk in client.chat_completion(
            model=model_id,
            messages=messages,
            max_tokens=8192,
            stream=True,
            temperature=0.7
        ):
            if chunk.choices[0].delta.content:
                response += chunk.choices[0].delta.content
                yield response
                
    except Exception as e:
        error_msg = str(e)
        if "401" in error_msg or "authorization" in error_msg.lower():
            yield "❌ Invalid API Key. Please check your Hugging Face token.\n\nGet a new one at: https://huggingface.co/settings/tokens"
        elif "429" in error_msg:
            yield "⏳ Rate limit exceeded. Please wait a moment and try again."
        else:
            yield f"❌ Error: {error_msg}\n\nPlease try again or select a different model."

CSS = """
@import url('https://fonts.googleapis.com/css2?family=JetBrains+Mono:wght@400;500;600;700&family=Space+Grotesk:wght@400;500;600;700&display=swap');

* {
    margin: 0;
    padding: 0;
    box-sizing: border-box;
}

:root {
    --bg-primary: #0A0E14;
    --bg-secondary: #0D1117;
    --bg-tertiary: #161B22;
    --border: #21262D;
    --text-primary: #E6EDF3;
    --text-secondary: #8B949E;
    --accent: #58A6FF;
    --accent-hover: #79C0FF;
    --success: #3FB950;
    --error: #F85149;
    --warning: #F0883E;
}

body {
    font-family: 'JetBrains Mono', monospace !important;
    background: var(--bg-primary) !important;
    color: var(--text-primary) !important;
}

.gradio-container {
    max-width: 100% !important;
    padding: 0 !important;
    margin: 0 !important;
    background: var(--bg-primary) !important;
}

#header {
    background: linear-gradient(135deg, var(--bg-secondary) 0%, var(--bg-tertiary) 100%);
    border-bottom: 2px solid var(--border);
    padding: 24px 48px;
    position: sticky;
    top: 0;
    z-index: 100;
}

#header h1 {
    font-family: 'Space Grotesk', sans-serif;
    font-size: 28px;
    font-weight: 700;
    letter-spacing: -0.02em;
    text-transform: uppercase;
    background: linear-gradient(135deg, var(--accent) 0%, var(--success) 100%);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
    margin: 0;
}

#header .subtitle {
    font-size: 12px;
    color: var(--text-secondary);
    margin-top: 4px;
    letter-spacing: 0.1em;
    text-transform: uppercase;
}

.sidebar-section {
    background: var(--bg-secondary);
    border: 1px solid var(--border);
    border-radius: 12px;
    padding: 20px;
    margin-bottom: 20px;
}

.sidebar-title {
    font-size: 11px;
    font-weight: 700;
    color: var(--text-secondary);
    text-transform: uppercase;
    letter-spacing: 0.1em;
    margin-bottom: 12px;
}

.api-key-input input {
    background: var(--bg-tertiary) !important;
    border: 1px solid var(--border) !important;
    color: var(--text-primary) !important;
    font-family: 'JetBrains Mono', monospace !important;
    font-size: 12px !important;
    padding: 12px !important;
    border-radius: 8px !important;
}

.api-key-input input:focus {
    border-color: var(--accent) !important;
    box-shadow: 0 0 0 3px rgba(88, 166, 255, 0.1) !important;
}

.security-note {
    background: rgba(63, 185, 80, 0.1);
    border: 1px solid var(--success);
    border-radius: 8px;
    padding: 12px;
    font-size: 11px;
    color: var(--success);
    line-height: 1.6;
    margin-top: 12px;
}

.warning-note {
    background: rgba(240, 136, 62, 0.1);
    border: 1px solid var(--warning);
    border-radius: 8px;
    padding: 12px;
    font-size: 11px;
    color: var(--warning);
    line-height: 1.6;
    margin-top: 12px;
}

.stat-grid {
    display: grid;
    grid-template-columns: 1fr 1fr;
    gap: 12px;
    margin-top: 16px;
}

.stat-box {
    background: var(--bg-tertiary);
    border: 1px solid var(--border);
    border-radius: 8px;
    padding: 12px;
    text-align: center;
}

.stat-value {
    font-size: 16px;
    font-weight: 700;
    color: var(--success);
    margin-bottom: 4px;
}

.stat-label {
    font-size: 9px;
    color: var(--text-secondary);
    text-transform: uppercase;
    letter-spacing: 0.1em;
}

.chatbot-container {
    background: transparent !important;
    border: none !important;
}

.message {
    margin-bottom: 16px !important;
}

.user .message-row {
    justify-content: flex-end !important;
}

.bot .message-row {
    justify-content: flex-start !important;
}

.message-content {
    padding: 16px 20px !important;
    border-radius: 16px !important;
    font-size: 14px !important;
    line-height: 1.6 !important;
    max-width: 70% !important;
}

.user .message-content {
    background: linear-gradient(135deg, var(--accent) 0%, var(--accent-hover) 100%) !important;
    color: var(--bg-primary) !important;
    border-radius: 16px 16px 4px 16px !important;
}

.bot .message-content {
    background: var(--bg-tertiary) !important;
    color: var(--text-primary) !important;
    border: 1px solid var(--border) !important;
    border-radius: 16px 16px 16px 4px !important;
}

#message-input textarea {
    background: var(--bg-tertiary) !important;
    border: 2px solid var(--border) !important;
    color: var(--text-primary) !important;
    font-family: 'JetBrains Mono', monospace !important;
    font-size: 14px !important;
    padding: 16px !important;
    border-radius: 12px !important;
}

#message-input textarea:focus {
    border-color: var(--accent) !important;
    box-shadow: 0 0 0 4px rgba(88, 166, 255, 0.1) !important;
}

.btn-primary {
    background: linear-gradient(135deg, var(--accent) 0%, var(--success) 100%) !important;
    color: var(--bg-primary) !important;
    font-family: 'Space Grotesk', sans-serif !important;
    font-weight: 600 !important;
    text-transform: uppercase !important;
    letter-spacing: 0.05em !important;
    border: none !important;
    padding: 14px 28px !important;
    border-radius: 8px !important;
    cursor: pointer !important;
}

.btn-primary:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 6px 20px rgba(88, 166, 255, 0.3) !important;
}

.btn-secondary {
    background: var(--bg-tertiary) !important;
    color: var(--text-primary) !important;
    font-family: 'Space Grotesk', sans-serif !important;
    font-weight: 600 !important;
    text-transform: uppercase !important;
    letter-spacing: 0.05em !important;
    border: 1px solid var(--border) !important;
    padding: 12px 24px !important;
    border-radius: 8px !important;
    cursor: pointer !important;
}

.btn-secondary:hover {
    background: var(--bg-secondary) !important;
    border-color: var(--accent) !important;
}

label {
    font-family: 'Space Grotesk', sans-serif !important;
    color: var(--text-secondary) !important;
    font-size: 11px !important;
    font-weight: 600 !important;
    text-transform: uppercase !important;
    letter-spacing: 0.05em !important;
}

.gr-dropdown {
    background: var(--bg-tertiary) !important;
    border: 1px solid var(--border) !important;
    color: var(--text-primary) !important;
    border-radius: 8px !important;
}

::selection {
    background: var(--accent);
    color: var(--bg-primary);
}

::-webkit-scrollbar {
    width: 10px;
}

::-webkit-scrollbar-track {
    background: var(--bg-secondary);
}

::-webkit-scrollbar-thumb {
    background: var(--border);
    border-radius: 5px;
}

::-webkit-scrollbar-thumb:hover {
    background: var(--text-secondary);
}
"""

with gr.Blocks(css=CSS, theme=gr.themes.Base(), title="Neural Chat") as demo:
    
    gr.HTML("""
    <div id="header">
        <h1>⬡ NEURAL CHAT</h1>
        <div class="subtitle">Secure AI Interface with Personal API Keys</div>
    </div>
    """)
    
    with gr.Row():
        
        with gr.Column(scale=1, min_width=320):
            
            with gr.Group(elem_classes=["sidebar-section"]):
                gr.HTML('<div class="sidebar-title">🔐 API AUTHENTICATION</div>')
                
                api_key_input = gr.Textbox(
                    label="Hugging Face API Key",
                    type="password",
                    placeholder="hf_...",
                    info="Your key is stored locally in your browser session",
                    elem_classes=["api-key-input"]
                )
                
                gr.HTML("""
                <div class="security-note">
                    ✓ Your API key is never stored on our servers<br>
                    ✓ Each user uses their own key<br>
                    ✓ Completely secure and private
                </div>
                """)
                
                gr.HTML("""
                <div class="warning-note">
                    📌 Get your free API key:<br>
                    <a href="https://huggingface.co/settings/tokens" target="_blank" style="color: #F0883E;">
                    https://huggingface.co/settings/tokens
                    </a>
                </div>
                """)
            
            with gr.Group(elem_classes=["sidebar-section"]):
                gr.HTML('<div class="sidebar-title">🤖 MODEL SELECTION</div>')
                
                model_selector = gr.Dropdown(
                    choices=[m["name"] for m in MODELS],
                    value=MODELS[0]["name"],
                    label="Active Model",
                    interactive=True,
                )
            
            with gr.Group(elem_classes=["sidebar-section"]):
                gr.HTML('<div class="sidebar-title">📊 STATISTICS</div>')
                
                gr.HTML("""
                <div class="stat-grid">
                    <div class="stat-box">
                        <div class="stat-value">●</div>
                        <div class="stat-label">SECURE</div>
                    </div>
                    <div class="stat-box">
                        <div class="stat-value">32K</div>
                        <div class="stat-label">CONTEXT</div>
                    </div>
                    <div class="stat-box">
                        <div class="stat-value">FREE</div>
                        <div class="stat-label">TIER</div>
                    </div>
                    <div class="stat-box">
                        <div class="stat-value">5</div>
                        <div class="stat-label">MODELS</div>
                    </div>
                </div>
                """)
            
            with gr.Group(elem_classes=["sidebar-section"]):
                gr.HTML('<div class="sidebar-title">⚡ ACTIONS</div>')
                
                clear_btn = gr.Button("⟳ NEW CHAT", elem_classes=["btn-secondary"], size="sm")
                retry_btn = gr.Button("🔄 RETRY LAST", elem_classes=["btn-secondary"], size="sm")
        
        with gr.Column(scale=3):
            
            chatbot = gr.Chatbot(
                value=[],
                height=600,
                show_label=False,
                type="messages",
                avatar_images=(None, "⬡"),
                elem_classes=["chatbot-container"]
            )
            
            msg_input = gr.Textbox(
                placeholder="Type your message... (Shift+Enter for new line)",
                show_label=False,
                lines=3,
                max_lines=10,
                elem_id="message-input"
            )
            
            with gr.Row():
                send_btn = gr.Button("SEND MESSAGE", elem_classes=["btn-primary"], scale=3)
                stop_btn = gr.Button("⏹ STOP", elem_classes=["btn-secondary"], scale=1)
            
            gr.HTML('<div class="sidebar-title" style="margin-top: 24px;">💡 EXAMPLE PROMPTS</div>')
            
            gr.Examples(
                examples=[
                    "Explain quantum computing in simple terms",
                    "Write a Python function for fibonacci sequence",
                    "What are the SOLID principles with examples?",
                    "Create a React component for a modal dialog",
                    "Help me optimize this SQL query",
                    "Explain async/await in JavaScript",
                ],
                inputs=msg_input,
            )
    
    def user_msg(message, history):
        if not message.strip():
            return "", history
        return "", history + [{"role": "user", "content": message}]
    
    def bot_msg(history, model, api_key):
        if not history or history[-1]["role"] != "user":
            return history
        
        user_message = history[-1]["content"]
        history.append({"role": "assistant", "content": ""})
        
        for response in stream_chat(user_message, history[:-1], model, api_key):
            history[-1]["content"] = response
            yield history
    
    msg_input.submit(
        user_msg,
        [msg_input, chatbot],
        [msg_input, chatbot],
        queue=False
    ).then(
        bot_msg,
        [chatbot, model_selector, api_key_input],
        chatbot
    )
    
    send_btn.click(
        user_msg,
        [msg_input, chatbot],
        [msg_input, chatbot],
        queue=False
    ).then(
        bot_msg,
        [chatbot, model_selector, api_key_input],
        chatbot
    )
    
    clear_btn.click(lambda: [], None, chatbot, queue=False)
    
    def retry_last_msg(history):
        if not history:
            return history
        
        if history[-1]["role"] == "assistant":
            history = history[:-1]
        
        if history and history[-1]["role"] == "user":
            last_user_msg = history[-1]
            history = history[:-1]
            return history + [last_user_msg]
        
        return history
    
    retry_btn.click(
        retry_last_msg,
        chatbot,
        chatbot,
        queue=False
    ).then(
        bot_msg,
        [chatbot, model_selector, api_key_input],
        chatbot
    )

if __name__ == "__main__":
    demo.queue(max_size=30)
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True
    )