import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import random
import warnings
import time
import threading
from typing import Dict, Any, Tuple

warnings.filterwarnings("ignore", category=UserWarning)

class AeroModel:
    def __init__(self, name: str, model_id: str, description: str, max_tokens: int, context_length: int, emoji: str):
        self.name = name
        self.model_id = model_id
        self.description = description
        self.max_tokens = max_tokens
        self.context_length = context_length
        self.emoji = emoji
        self.model = None
        self.tokenizer = None
        self.loaded = False
        self.cache = {}
    
    def load_model(self):
        if self.loaded:
            return
        
        print(f"🔄 Loading {self.name}...")
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_id, trust_remote_code=True)
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token
        
        self.model = AutoModelForCausalLM.from_pretrained(
            self.model_id,
            torch_dtype=torch.float32,
            trust_remote_code=True,
            low_cpu_mem_usage=False
        )
        self.model.eval()
        self.loaded = True
        print(f"✅ {self.name} loaded successfully!")
    
    def generate_response(self, prompt: str, history: list) -> str:
        if not self.loaded:
            return f"❌ {self.name} not loaded. Please wait for initialization."
        
        # Check cache for similar responses
        cache_key = prompt.lower().strip()
        if cache_key in self.cache:
            return f"⚡ {self.cache[cache_key]}"
        
        # Build context with history
        context = ""
        for user_msg, bot_msg in history[-3:]:  # Keep last 3 exchanges
            context += f"Human: {user_msg}\nAero: {bot_msg}\n"
        context += f"Human: {prompt}\nAero:"
        
        # Tokenize and generate
        inputs = self.tokenizer.encode(context, return_tensors="pt", max_length=self.context_length, truncation=True)
        
        with torch.no_grad():
            outputs = self.model.generate(
                inputs,
                max_new_tokens=self.max_tokens,
                min_new_tokens=10,
                temperature=0.8,
                top_k=50,
                top_p=0.9,
                do_sample=True,
                pad_token_id=self.tokenizer.eos_token_id,
                repetition_penalty=1.1
            )
        
        response = self.tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)
        response = response.strip()
        
        # Cache good responses
        if len(response) > 20 and len(self.cache) < 50:
            self.cache[cache_key] = response
        
        return response

# Define the Aero model lineup
AERO_MODELS = {
    "basic": AeroModel(
        name="Aero Basic",
        model_id="gpt2",
        description="Fast and reliable - perfect for everyday conversations",
        max_tokens=60,
        context_length=512,
        emoji="🔵"
    ),
    "smart": AeroModel(
        name="Aero Smart",
        model_id="microsoft/phi-2",
        description="Advanced reasoning and complex problem solving (slower)",
        max_tokens=150,
        context_length=768,
        emoji="🧠"
    ),
    "speed": AeroModel(
        name="Aero Speed",
        model_id="distilgpt2",
        description="Lightning-fast responses for quick interactions",
        max_tokens=40,
        context_length=256,
        emoji="⚡"
    )
}

class AeroAISystem:
    def __init__(self):
        self.current_model = "basic"
        self.loading_models = set()
    
    def load_model_async(self, model_key: str):
        if model_key in self.loading_models or AERO_MODELS[model_key].loaded:
            return
        
        self.loading_models.add(model_key)
        thread = threading.Thread(target=self._load_model_thread, args=(model_key,))
        thread.daemon = True
        thread.start()
    
    def _load_model_thread(self, model_key: str):
        try:
            AERO_MODELS[model_key].load_model()
        except Exception as e:
            print(f"❌ Failed to load {AERO_MODELS[model_key].name}: {e}")
        finally:
            self.loading_models.discard(model_key)
    
    def switch_model(self, model_choice: str) -> Tuple[str, str]:
        model_map = {
            "🔵 Aero Basic - Fast & Reliable": "basic",
            "🧠 Aero Smart - Advanced Reasoning": "smart", 
            "⚡ Aero Speed - Lightning Fast": "speed"
        }
        
        model_key = model_map.get(model_choice, "basic")
        self.current_model = model_key
        
        # Load model if not already loaded
        if not AERO_MODELS[model_key].loaded:
            self.load_model_async(model_key)
        
        model = AERO_MODELS[model_key]
        status = f"🔄 Switching to {model.name} {model.emoji}"
        info = f"**{model.name}** - {model.description}\n\nMax tokens: {model.max_tokens} | Context: {model.context_length}"
        
        return status, info
    
    def chat_with_aero(self, message: str, history: list) -> Tuple[list, str]:
        if not message.strip():
            return history, ""
        
        model = AERO_MODELS[self.current_model]
        
        if not model.loaded:
            if self.current_model in self.loading_models:
                bot_response = f"🔄 {model.name} is still loading... Please wait a moment."
            else:
                self.load_model_async(self.current_model)
                bot_response = f"🔄 Loading {model.name}... This may take a moment for the first time."
        else:
            try:
                bot_response = model.generate_response(message, history)
            except Exception as e:
                bot_response = f"❌ Error: {str(e)}"
        
        history.append([message, bot_response])
        return history, ""

# Initialize the system
aero_system = AeroAISystem()

def create_interface():
    # Startup banner
    print("""
╔═══════════════════════════════════════════════════════════════════════════════╗
║                                                                               ║
║     ░█████╗░███████╗██████╗░░█████╗░░█████╗░██╗                              ║
║     ██╔══██╗██╔════╝██╔══██╗██╔══██╗██╔══██╗██║                              ║
║     ███████║█████╗░░██████╔╝██║░░██║███████║██║                              ║
║     ██╔══██║██╔══╝░░██╔══██╗██║░░██║██╔══██║██║                              ║
║     ██║░░██║███████╗██║░░██║╚█████╔╝██║░░██║██║                              ║
║     ╚═╝░░╚═╝╚══════╝╚═╝░░╚═╝░╚════╝░╚═╝░░╚═╝╚═╝                              ║
║                                                                               ║
║                🚀 MULTI-MODEL NEURAL INTERFACE 🚀                            ║
║                    B L A C K L I N K  L A B S                               ║
╚═══════════════════════════════════════════════════════════════════════════════╝
    """)
    
    # Load Aero Basic by default
    aero_system.load_model_async("basic")
    
    # Custom CSS for the interface
    css = """
    .gradio-container {
        background: linear-gradient(135deg, #0c0c0c 0%, #1a1a2e 50%, #16213e 100%);
        color: #00ff88;
        font-family: 'Courier New', monospace;
    }
    
    .model-selector {
        background: rgba(0, 255, 136, 0.1);
        border: 1px solid #00ff88;
        border-radius: 8px;
        padding: 10px;
    }
    
    .status-display {
        background: rgba(0, 100, 255, 0.1);
        border: 1px solid #0064ff;
        border-radius: 8px;
        padding: 10px;
        color: #0064ff;
    }
    
    .chat-container {
        background: rgba(0, 0, 0, 0.3);
        border: 1px solid #00ff88;
        border-radius: 10px;
    }
    
    .input-box {
        background: rgba(0, 255, 136, 0.05);
        border: 2px solid #00ff88;
        color: #00ff88;
    }
    
    .aero-button {
        background: linear-gradient(45deg, #00ff88, #0064ff);
        border: none;
        color: black;
        font-weight: bold;
        transition: all 0.3s ease;
    }
    
    .aero-button:hover {
        transform: scale(1.05);
        box-shadow: 0 0 20px rgba(0, 255, 136, 0.5);
    }
    """
    
    with gr.Blocks(css=css, title="AeroAI Multi-Model System") as interface:
        
        gr.HTML("""
        <div style="text-align: center; padding: 20px;">
            <h1 style="color: #00ff88; text-shadow: 0 0 10px #00ff88;">🚀 AeroAI Multi-Model System</h1>
            <p style="color: #0064ff;">Choose your AI model based on your needs</p>
        </div>
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                model_selector = gr.Dropdown(
                    choices=[
                        "🔵 Aero Basic - Fast & Reliable",
                        "🧠 Aero Smart - Advanced Reasoning", 
                        "⚡ Aero Speed - Lightning Fast"
                    ],
                    value="🔵 Aero Basic - Fast & Reliable",
                    label="🤖 Select Aero Model",
                    elem_classes=["model-selector"]
                )
                
                model_status = gr.Textbox(
                    value="🔵 Aero Basic - Ready to chat!",
                    label="🔋 Model Status",
                    interactive=False,
                    elem_classes=["status-display"]
                )
                
                model_info = gr.Markdown(
                    value="**Aero Basic** - Fast and reliable - perfect for everyday conversations\n\nMax tokens: 60 | Context: 512",
                    elem_classes=["status-display"]
                )
            
            with gr.Column(scale=2):
                chatbot = gr.Chatbot(
                    label="💬 Chat with AeroAI",
                    height=400,
                    elem_classes=["chat-container"]
                )
                
                msg = gr.Textbox(
                    placeholder="Type your message here... 🚀",
                    label="Message",
                    elem_classes=["input-box"]
                )
                
                with gr.Row():
                    send_btn = gr.Button("Send 🚀", elem_classes=["aero-button"])
                    clear_btn = gr.Button("Clear Chat 🗑️", elem_classes=["aero-button"])
        
        # Example prompts
        gr.HTML("""
        <div style="text-align: center; margin-top: 20px; color: #00ff88;">
            <p><strong>Try these prompts:</strong></p>
            <p>💡 "Explain quantum computing" | 🎨 "Write a short story" | 💻 "Help me debug this code"</p>
        </div>
        """)
        
        # Event handlers
        model_selector.change(
            fn=aero_system.switch_model,
            inputs=[model_selector],
            outputs=[model_status, model_info]
        )
        
        send_btn.click(
            fn=aero_system.chat_with_aero,
            inputs=[msg, chatbot],
            outputs=[chatbot, msg]
        )
        
        msg.submit(
            fn=aero_system.chat_with_aero,
            inputs=[msg, chatbot], 
            outputs=[chatbot, msg]
        )
        
        clear_btn.click(
            fn=lambda: ([], ""),
            outputs=[chatbot, msg]
        )
    
    return interface

if __name__ == "__main__":
    print("🌟 Initializing AeroAI Multi-Model System...")
    print("🔋 System Status: Online")
    print("⚡ Ready for neural connection\n")
    
    demo = create_interface()
    demo.launch(
        share=True,
        inbrowser=False,
        show_error=True
    )