import os
import json
from typing import Optional
from datetime import datetime
import google.generativeai as genai
import gradio as gr
import bcrypt
from tinydb import TinyDB, Query

GEMINI_API_KEY = "AIzaSyBSHcHdbVcvo5IvKDAq2fY1xsWDDb2-JUI"
MODEL_NAME     = "models/gemini-2.5-flash"
DB_PATH        = "users_db.json"

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

# Initialize TinyDB
db = TinyDB(DB_PATH)
User = Query()

# Global state for current user session
current_user = {"logged_in": False, "user_data": None}

class UserAuthentication:
    """Handle user authentication and database operations"""
    
    def __init__(self, db_path: str = DB_PATH):
        self.db = TinyDB(db_path)
    
    @staticmethod
    def hash_password(password: str) -> str:
        """Hash password using bcrypt"""
        return bcrypt.hashpw(password.encode('utf-8'), bcrypt.gensalt()).decode('utf-8')
    
    @staticmethod
    def verify_password(password: str, hashed: str) -> bool:
        """Verify password against hash"""
        return bcrypt.checkpw(password.encode('utf-8'), hashed.encode('utf-8'))
    
    def user_exists(self, email: str) -> bool:
        """Check if user already exists"""
        return self.db.search(User.email == email) != []
    
    def signup(self, username: str, email: str, password: str, phone: str = "") -> tuple[bool, str]:
        """Register a new user"""
        try:
            if not username or not email or not password:
                return False, "All fields are required"
            
            if self.user_exists(email):
                return False, "Email already registered"
            
            if len(password) < 6:
                return False, "Password must be at least 6 characters"
            
            hashed_password = self.hash_password(password)
            
            user_data = {
                "username": username,
                "email": email,
                "password": hashed_password,
                "phone": phone,
                "created_at": datetime.now().isoformat(),
                "last_login": None
            }
            
            self.db.insert(user_data)
            return True, f"✓ Account created successfully! Welcome, {username}!"
        
        except Exception as e:
            return False, f"Signup error: {str(e)}"
    
    def login(self, email: str, password: str) -> tuple[bool, str, Optional[dict]]:
        """Authenticate user"""
        try:
            if not email or not password:
                return False, "Email and password required", None
            
            users = self.db.search(User.email == email)
            
            if not users:
                return False, "Email not found", None
            
            user = users[0]
            
            if not self.verify_password(password, user['password']):
                return False, "Incorrect password", None
            
            # Update last login
            self.db.update(
                {'last_login': datetime.now().isoformat()},
                User.email == email
            )
            
            return True, f"✓ Login successful! Welcome, {user['username']}!", user
        
        except Exception as e:
            return False, f"Login error: {str(e)}", None
    
    def get_user_profile(self, email: str) -> Optional[dict]:
        """Get user profile"""
        users = self.db.search(User.email == email)
        if users:
            user = users[0].copy()
            user.pop('password', None)  # Remove password from profile
            return user
        return None

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 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)

# Initialize authentication system
auth_system = UserAuthentication(DB_PATH)

# Create demo account if it doesn't exist
demo_email = "demo@example.com"
if not auth_system.user_exists(demo_email):
    auth_system.signup("Demo User", demo_email, "password123")
    print("✅ Demo account created successfully!")
    print("   Email: demo@example.com")
    print("   Password: password123")

# ===== AUTHENTICATION FUNCTIONS =====
def signup_user(username: str, email: str, password: str, confirm_password: str) -> tuple[str, str]:
    """Handle user signup"""
    if not password or password != confirm_password:
        return "❌ Passwords do not match!", ""
    
    success, message = auth_system.signup(username, email, password)
    if success:
        return "✅ " + message, ""
    else:
        return "❌ " + message, ""

def login_user(email: str, password: str) -> tuple[str, str, bool]:
    """Handle user login"""
    success, message, user = auth_system.login(email, password)
    if success:
        current_user["logged_in"] = True
        current_user["user_data"] = user
        username = user['username'] if user else ""
        return "✅ " + message, username, True
    else:
        return "❌ " + message, "", False

def logout_user() -> None:
    """Handle user logout"""
    current_user["logged_in"] = False
    current_user["user_data"] = None
    chatbot.conversation_history = []

# ===== CHAT FUNCTIONS WITH USER SESSION =====
def respond_authenticated(message: str, chat_history: list, mode: str, temperature: float):
    """Respond to user message (only when authenticated)"""
    if not current_user["logged_in"]:
        return "", 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 get_current_username() -> str:
    """Get currently logged-in username"""
    if current_user["logged_in"] and current_user["user_data"]:
        return current_user["user_data"].get("username", "User")
    return "Guest"

# Create Unified Gradio Application
with gr.Blocks(title="Nexus ChatBot", theme=gr.themes.Soft()) as demo:
    # State to track if user is logged in
    auth_state = gr.State(value=False)
    
    gr.Markdown("""
    # 🤖 Nexus Intelligent ChatBot
    
    Your Intelligent AI Assistant powered by Google's Gemini AI
    """)
    
    # AUTHENTICATION SECTION (shown when not logged in)
    with gr.Group(visible=True) as auth_section:
        gr.Markdown("### 🔐 Authentication")
        
        with gr.Tabs():
            # LOGIN TAB
            with gr.TabItem("🔐 Login"):
                gr.Markdown("#### Sign In to Your Account")
                
                with gr.Column():
                    login_email = gr.Textbox(
                        label="Email",
                        placeholder="Enter your email",
                        type="text"
                    )
                    login_password = gr.Textbox(
                        label="Password",
                        placeholder="Enter your password",
                        type="password"
                    )
                    login_msg = gr.Textbox(
                        label="Status",
                        interactive=False,
                        visible=False
                    )
                    
                    login_btn = gr.Button("🔓 Login", variant="primary", size="lg")
                    
                    # Success notification
                    login_popup = gr.HTML(
                        value="",
                        visible=False
                    )
                    
                    gr.Markdown("---")
                    gr.Markdown("""
                    **📌 Demo Account Available:**
                    - Email: `demo@example.com`
                    - Password: `password123`
                    
                    Or create a new account in the Sign Up tab!
                    """)
            
            # SIGNUP TAB
            with gr.TabItem("📝 Sign Up"):
                gr.Markdown("#### Create a New Account")
                
                with gr.Column():
                    signup_username = gr.Textbox(
                        label="Username",
                        placeholder="Choose a username",
                        type="text"
                    )
                    signup_email = gr.Textbox(
                        label="Email",
                        placeholder="Enter your email",
                        type="text"
                    )
                    signup_password = gr.Textbox(
                        label="Password",
                        placeholder="Create a password (min 6 chars)",
                        type="password"
                    )
                    signup_confirm = gr.Textbox(
                        label="Confirm Password",
                        placeholder="Confirm your password",
                        type="password"
                    )
                    signup_msg = gr.Textbox(
                        label="✅ Status",
                        interactive=False,
                        visible=False
                    )
                    
                    signup_btn = gr.Button("✍️ Sign Up", variant="primary", size="lg")
                    
                    # Success notification
                    success_popup = gr.HTML(
                        value="",
                        visible=False
                    )
    
    # CHAT SECTION (shown when logged in)
    with gr.Group(visible=False) as chat_section:
        # User info bar
        with gr.Row():
            with gr.Column(scale=4):
                gr.Markdown("### 💬 Chat Interface")
            with gr.Column(scale=1):
                current_user_display = gr.Textbox(
                    label="👤 Logged In As",
                    value=get_current_username(),
                    interactive=False
                )
        
        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="Creativity",
                    info="Higher = more creative, Lower = 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", scale=1)
            logout_btn = gr.Button("🚪 Logout", variant="stop", scale=1)
        
        export_output = gr.Textbox(
            label="📋 Exported Chat (JSON)",
            interactive=False,
            visible=False
        )
        
        gr.Markdown("""
        ### 📝 Chat Modes:
        - **General**: Everyday questions and assistance
        - **Technical**: Code help and tech support
        - **Creative**: Storytelling and creative writing
        - **Educational**: Learning and explanations
        - **Medical**: Health information
        """)
    
    # === EVENT HANDLERS ===
    
    # Login handler
    def handle_login(email: str, password: str):
        success, message, user = auth_system.login(email, password)
        if success:
            current_user["logged_in"] = True
            current_user["user_data"] = user
            username = user['username'] if user else ""
            success_html = f'<div style="background-color: #ccffcc; border: 2px solid #00cc00; padding: 15px; border-radius: 5px; color: #009900; font-weight: bold; text-align: center; font-size: 18px;">✅ {message}<br><small>Redirecting to chat...</small></div>'
            return (
                success_html,
                gr.update(visible=False),  # Hide auth section
                gr.update(visible=True),   # Show chat section
                username
            )
        else:
            error_html = f'<div style="background-color: #ffcccc; border: 2px solid #ff0000; padding: 15px; border-radius: 5px; color: #cc0000; font-weight: bold;">❌ {message}</div>'
            return error_html, gr.update(visible=True), gr.update(visible=False), ""
    
    # Signup handler
    def handle_signup(username: str, email: str, password: str, confirm_password: str):
        if not password or password != confirm_password:
            error_html = '<div style="background-color: #ffcccc; border: 2px solid #ff0000; padding: 15px; border-radius: 5px; color: #cc0000; font-weight: bold;">❌ Passwords do not match!</div>'
            return "❌ Passwords do not match!", error_html, gr.update(visible=True)
        
        success, message = auth_system.signup(username, email, password)
        if success:
            success_html = f'<div style="background-color: #ccffcc; border: 2px solid #00cc00; padding: 15px; border-radius: 5px; color: #009900; font-weight: bold; text-align: center; font-size: 18px;">✅ {message}<br><small>You can now login with your email and password</small></div>'
            return message, success_html, gr.update(visible=True)
        else:
            error_html = f'<div style="background-color: #ffcccc; border: 2px solid #ff0000; padding: 15px; border-radius: 5px; color: #cc0000; font-weight: bold;">❌ {message}</div>'
            return message, error_html, gr.update(visible=True)
    
    # Logout handler
    def handle_logout():
        current_user["logged_in"] = False
        current_user["user_data"] = None
        chatbot.conversation_history = []
        return (
            gr.update(visible=True),   # Show auth section
            gr.update(visible=False),  # Hide chat section
            "",                         # Clear message
            [],                         # Clear chat history
        )
    
    # Login click event
    login_btn.click(
        handle_login,
        inputs=[login_email, login_password],
        outputs=[login_popup, auth_section, chat_section, current_user_display]
    )
    
    # Signup click event
    signup_btn.click(
        handle_signup,
        inputs=[signup_username, signup_email, signup_password, signup_confirm],
        outputs=[signup_msg, success_popup, success_popup]
    )
    
    # Chat send button
    send_btn.click(
        respond_authenticated,
        inputs=[msg_input, chatbot_ui, mode, temperature],
        outputs=[msg_input, chatbot_ui]
    )
    
    # Chat input submit (Enter key)
    msg_input.submit(
        respond_authenticated,
        inputs=[msg_input, chatbot_ui, mode, temperature],
        outputs=[msg_input, chatbot_ui]
    )
    
    # Clear chat
    clear_btn.click(
        clear_history,
        outputs=[chatbot_ui, msg_input]
    )
    
    # Export chat
    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]
    )
    
    # Logout
    logout_btn.click(
        handle_logout,
        outputs=[auth_section, chat_section, login_popup, chatbot_ui]
    )

# Main application logic
if __name__ == "__main__":
    """
    Start the Gradio application
    - Authentication page (Login/Sign Up) appears first
    - After successful login, chat interface opens automatically
    - Logout returns to authentication
    """
    print("🚀 Starting Nexus ChatBot Application...")
    print("📍 Opening Authentication Page...")
    print("\n💡 Demo Account for Testing:")
    print("   Email: demo@example.com")
    print("   Password: password123")
    print("\n✨ Features:")
    print("   • Secure login and registration")
    print("   • Multiple chat modes (general, technical, creative, educational, medical)")
    print("   • Export conversations")
    print("   • Powered by Google's Gemini AI")
    
    demo.launch(share=False)