import os
import json
from typing import Optional
from datetime import datetime
import google.generativeai as genai
import gradio as gr

GEMINI_API_KEY = "AIzaSyBKlMnddESWNPvvEJIfu5_qum6wKHm-EnU"
MODEL_NAME     = "models/gemini-2.5-flash"

# Configure the API
genai.configure(api_key=GEMINI_API_KEY)

class GeminiChatBot:
    """Main chatbot class with context management and multiple modes"""
    
    def __init__(self, model_name: str = MODEL_NAME):
        self.model = genai.GenerativeModel(model_name)
        self.conversation_history = []
        self.chat_session = None
        self.system_prompt = ""
        
    def set_system_prompt(self, mode: str):
        """Set system prompt based on chatbot mode"""
        prompts = {
            "general": """You are a helpful, accurate, and friendly AI assistant. 
            Provide clear, concise, and informative responses. 
            Always be honest about limitations and uncertainty.""",
            
            "technical": """You are an expert technical support assistant. 
            Provide detailed technical solutions, code examples, and best practices. 
            When unsure, ask clarifying questions. Always suggest verification steps.""",
            
            "creative": """You are a creative writing assistant with strong storytelling abilities. 
            Help users with creative writing, brainstorming, and narrative development. 
            Provide engaging and imaginative content.""",
            
            "educational": """You are an educational tutor. Explain concepts clearly, 
            break down complex topics, and provide examples. 
            Encourage learning and ask clarifying questions.""",
            
            "medical": """You are a medical information assistant. 
            Provide accurate health information and general guidance. 
            Always recommend consulting healthcare professionals for serious concerns. 
            Do NOT provide emergency medical advice."""
        }
        self.system_prompt = prompts.get(mode, prompts["general"])
    
    def chat(self, user_message: str, mode: str = "general", temperature: float = 0.7) -> str:
        """Generate response using Gemini with context"""
        try:
            self.set_system_prompt(mode)
            
            # Build messages with system context
            messages = [
                {"role": "user", "parts": [f"[SYSTEM: {self.system_prompt}]\n\n{user_message}"]}
            ]
            
            # Add conversation history for context (last 5 exchanges)
            if self.conversation_history:
                history_context = "\n".join([
                    f"Previous: {msg}" 
                    for msg in self.conversation_history[-5:]
                ])
                full_message = f"[Conversation Context]\n{history_context}\n\n[New Message]\n{user_message}"
            else:
                full_message = user_message
            
            # Generate response
            response = self.model.generate_content(
                full_message,
                generation_config=genai.types.GenerationConfig(
                    temperature=temperature,
                    top_p=0.95,
                    top_k=40
                )
            )
            
            bot_response = response.text
            
            # Store in history
            self.conversation_history.append(f"User: {user_message[:100]}...")
            self.conversation_history.append(f"Bot: {bot_response[:100]}...")
            
            return bot_response
            
        except Exception as e:
            return f"Error: {str(e)}\n\nMake sure your API key is valid. Get it from: https://aistudio.google.com/app/apikey"

# Initialize chatbot
chatbot = GeminiChatBot()

# Gradio Interface Functions
def respond(message: str, chat_history: list, mode: str, temperature: float):
    """Respond to user message and return updated chat history"""
    response = chatbot.chat(message, mode=mode, temperature=temperature)
    chat_history.append({"role": "user", "content": message})
    chat_history.append({"role": "assistant", "content": response})
    return "", chat_history

def clear_history():
    """Clear conversation history"""
    chatbot.conversation_history = []
    return [], ""

def export_chat(chat_history: list) -> str:
    """Export chat as JSON"""
    if not chat_history:
        return "No chat history to export"
    
    export_data = {
        "timestamp": datetime.now().isoformat(),
        "conversation": chat_history
    }
    return json.dumps(export_data, indent=2)

# Create Gradio Interface
with gr.Blocks(title="Gemini ChatBot", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🤖 Nexus Intelligent ChatBot
    
    A generalized, accurate chatbot powered by Google's Gemini AI.
    Select your mode and start chatting!
    """)
    
    with gr.Row():
        with gr.Column(scale=3):
            chatbot_ui = gr.Chatbot(
                label="Chat History",
                height=500,
                show_label=True
            )
        
        with gr.Column(scale=1):
            gr.Markdown("### ⚙️ Settings")
            
            mode = gr.Radio(
                choices=["general", "technical", "creative", "educational", "medical"],
                value="general",
                label="Chat Mode",
                info="Select conversation style"
            )
            
            temperature = gr.Slider(
                minimum=0,
                maximum=2,
                value=0.7,
                step=0.1,
                label="Temperature",
                info="Higher = more creative, Lower = more focused"
            )
    
    with gr.Row():
        msg_input = gr.Textbox(
            placeholder="Type your message here...",
            label="Your Message",
            lines=2
        )
    
    with gr.Row():
        send_btn = gr.Button("Send", variant="primary", scale=2)
        clear_btn = gr.Button("Clear Chat", scale=1)
        export_btn = gr.Button("Export Chat", scale=1)
    
    export_output = gr.Textbox(
        label="Exported Chat (JSON)",
        interactive=False,
        visible=False
    )
    
    # Event handlers
    send_btn.click(
        respond,
        inputs=[msg_input, chatbot_ui, mode, temperature],
        outputs=[msg_input, chatbot_ui]
    )
    
    msg_input.submit(
        respond,
        inputs=[msg_input, chatbot_ui, mode, temperature],
        outputs=[msg_input, chatbot_ui]
    )
    
    clear_btn.click(
        clear_history,
        outputs=[chatbot_ui, msg_input]
    )
    
    def toggle_export_visibility():
        return gr.update(visible=True)
    
    def get_and_show_export(chat_history):
        return export_chat(chat_history), gr.update(visible=True)
    
    export_btn.click(
        get_and_show_export,
        inputs=[chatbot_ui],
        outputs=[export_output, export_output]
    )
    
    gr.Markdown("""
    ### 📝 Chat Modes:
    - **General**: Friendly assistant for everyday questions
    - **Technical**: Expert technical support and code help
    - **Creative**: Storytelling and creative writing
    - **Educational**: Learning and concept explanation
    - **Medical**: Health information (consult professionals for serious concerns)
    
    ### 🔑 Setup:
    1. Get your API key from [Google AI Studio](https://aistudio.google.com/app/apikey)
    2. Set `GEMINI_API_KEY` in `.env` file or environment variables
    3. Run the app and start chatting!
    """)

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