app.py

from fastapi import FastAPI, HTTPException, Header, Response
from fastapi.responses import HTMLResponse, StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import google.generativeai as genai
import httpx
import os
from dotenv import load_dotenv
from duckduckgo_search import DDGS
from typing import Optional, List, Dict, Any
import logging
import re
import asyncio
import threading
import time
import hashlib
import json
from functools import wraps
from collections import OrderedDict
from dataclasses import dataclass, field

# Load environment variables from .env file
load_dotenv()


import spacy
from rake_nltk import Rake
import nltk
from cache.chromadb_cache import ChromaDBSearchCache
from search_optimizer import (
    has_meaningful_conversation_history,
    hybrid_search_decision,
    extract_search_terms,
    format_search_context
)

nltk.download('stopwords')
nltk.download('punkt_tab')

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Initialize FastAPI app
app = FastAPI(
    title="Enhanced Chat API with Dynamic Model Selection",
    description="API with query classification and dynamic model loading for QA/Summarization",
    version="3.0.0"
)

# Configure CORS with restricted origins for security
app.add_middleware(
    CORSMiddleware,
    allow_origins=[
        "https://huggingface.co",
        "https://*.hf.space",
        "http://localhost:3000",
        "http://localhost:8000",
        "http://127.0.0.1:3000",
        "http://127.0.0.1:8000"
    ],
    allow_methods=["POST", "GET"],
    allow_headers=["Content-Type", "Authorization"],
)

# Request/Response models
class ChatRequest(BaseModel):
    prompt: str
    max_new_tokens: int = 500
    use_search: bool = True
    temperature: float = 0.7
    user_id: Optional[str] = None
    history: Optional[List[Dict[str, str]]] = None
    force_search: Optional[bool] = None
    search_decision_mode: str = "balanced"  # conservative, balanced, aggressive

class ChatResponse(BaseModel):
    response: str
    search_results: Optional[List[Dict[str, Any]]] = None
    search_decision: Optional[Dict[str, Any]] = None
    cache_info: Optional[Dict[str, Any]] = None

class SearchRequest(BaseModel):
    query: str
    max_results: int = 5


# Configure Google AI
try:
    genai.configure(api_key=os.environ["GOOGLE_API_KEY"])
except KeyError:
    logger.error("CRITICAL: GOOGLE_API_KEY environment variable not set.")
    raise # Exit if API key is not set

# Initialize the Generative Model
model = genai.GenerativeModel('gemini-1.5-flash')

# Global NLP tools
nlp = spacy.load("en_core_web_sm")
rake = Rake()

# ===== Phase 3c: ChromaDB Vector Database Cache System =====

# Universal ChromaDB search cache - single global instance for all users
# Vector database provides better semantic matching and persistent storage
universal_search_cache = None

def get_universal_cache() -> ChromaDBSearchCache:
    """Get the universal ChromaDB search cache instance"""
    global universal_search_cache
    if universal_search_cache is None:
        # Initialize ChromaDB cache with environment configuration
        cache_db_path = os.getenv('CHROMADB_PATH', 'cache_db')
        cache_results_path = os.getenv('CACHE_RESULTS_PATH', 'cache_results')
        embedding_model = os.getenv('CACHE_EMBEDDING_MODEL', 'all-MiniLM-L6-v2')
        
        universal_search_cache = ChromaDBSearchCache(
            max_size=1000,  # Large cache size
            default_ttl=3600,  # 1 hour TTL
            cache_db_path=cache_db_path,
            cache_results_path=cache_results_path,
            embedding_model=embedding_model,
            similarity_threshold=0.7
        )
        logger.info(f"Initialized ChromaDB universal cache: {cache_db_path}")
    
    return universal_search_cache

# ===== End Phase 3c ChromaDB Cache System =====




def normalize_user_id(user_id: Optional[str]) -> Optional[str]:
    """Normalize user_id to None for anonymous requests"""
    if user_id is None or user_id.strip() == "":
        return None
    return user_id.strip()

def validate_user_id(user_id: Optional[str]) -> Optional[str]:
    """Validate user_id format and return normalized value"""
    # Normalize to None for anonymous users
    user_id = normalize_user_id(user_id)
    
    if user_id is None:
        return None  # Anonymous user
    
    # Validate format: non-empty string, max 255 chars, alphanumeric + hyphens + underscores
    if len(user_id) > 255:
        raise HTTPException(status_code=400, detail="user_id must be 255 characters or less")
    
    # Check allowed characters
    if not re.match(r'^[a-zA-Z0-9_-]+$', user_id):
        raise HTTPException(status_code=400, detail="user_id can only contain alphanumeric characters, hyphens, and underscores")
    
    return user_id

def run_in_threadpool(func):
    """Decorator to run synchronous model inference in thread pool"""
    @wraps(func)
    async def wrapper(*args, **kwargs):
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(None, func, *args, **kwargs)
    return wrapper

async def search_brave(query: str, max_results: int = 5) -> List[Dict[str, Any]]:
    """Search using Brave Search API with async HTTP client"""
    try:
        # Check for API key in environment variable first, then fallback to hardcoded
        api_key = os.getenv('BRAVE_API_KEY')
        
        if not api_key:
            logger.error("No Brave API key available")
            return []
        
        headers = {
            'Accept': 'application/json',
            'Accept-Encoding': 'gzip',
            'X-Subscription-Token': api_key
        }
        
        params = {
            'q': query,
            'count': max_results,
            'safesearch': 'moderate',
            'search_lang': 'en',
            'country': 'US'
        }
        
        
        async with httpx.AsyncClient(timeout=10.0) as client:
            response = await client.get(
                'https://api.search.brave.com/res/v1/web/search',
                headers=headers,
                params=params
            )
            
            response.raise_for_status()
            
            data = response.json()
            
            web_results = data.get('web', {})
            raw_results = web_results.get('results', [])
            
            results = []
            for result in raw_results:
                results.append({
                    "title": result.get("title", ""),
                    "body": re.sub(r'\s+', ' ', result.get("description", "")).strip(),
                    "href": result.get("url", ""),
                    "source": "Brave"
                })
            
            return results[:max_results]
        
    except Exception as e:
        logger.error(f"Brave Search error: {e}")
        logger.error(f"Brave Search error type: {type(e).__name__}")
        import traceback
        logger.error(f"Brave Search traceback: {traceback.format_exc()}")
        return []

async def search_duckduckgo(query: str, max_results: int = 5) -> List[Dict[str, Any]]:
    """Search using DuckDuckGo with async execution and timeout handling"""
    try:
        # Run DuckDuckGo search in thread pool with timeout
        loop = asyncio.get_event_loop()
        results = await asyncio.wait_for(
            loop.run_in_executor(None, _sync_duckduckgo_search, query, max_results),
            timeout=8.0  # 8 second timeout for Hugging Face compatibility
        )
        return results
    except asyncio.TimeoutError:
        logger.warning(f"DuckDuckGo search timed out for query: {query}")
        return []
    except Exception as e:
        logger.error(f"DuckDuckGo Search error: {e}")
        return []

def _sync_duckduckgo_search(query: str, max_results: int) -> List[Dict[str, Any]]:
    """Synchronous DuckDuckGo search helper with retry logic"""
    max_retries = 2
    for attempt in range(max_retries):
        try:
            # Configure DDGS with more conservative settings for hosted environments
            with DDGS(timeout=5) as ddgs:
                results = []
                search_results = ddgs.text(
                    query, 
                    safesearch='moderate', 
                    max_results=max_results,
                    region='us-en'  # Specify region to potentially avoid some blocks
                )
                
                for result in search_results:
                    results.append({
                        "title": result.get("title", ""),
                        "body": re.sub(r'\s+', ' ', result.get("body", "")).strip(),
                        "href": result.get("href", ""),
                        "source": "DuckDuckGo"
                    })
                
                logger.info(f"DuckDuckGo search successful on attempt {attempt + 1}")
                return results
                
        except Exception as e:
            logger.warning(f"DuckDuckGo attempt {attempt + 1} failed: {e}")
            if attempt == max_retries - 1:  # Last attempt
                logger.error(f"DuckDuckGo search failed after {max_retries} attempts")
                raise
            # Wait briefly before retry
            import time
            time.sleep(1)
    
    return []

async def search_web_combined(query: str, max_results: int = 10) -> List[Dict[str, Any]]:
    """Combined web search with resilient fallback strategy"""
    try:
        logger.info(f"Combined Search: Starting search for '{query}'")
        
        # Run both searches concurrently with timeout protection
        brave_task = asyncio.create_task(search_brave(query, 6))  # Get more from Brave as primary
        duckduckgo_task = asyncio.create_task(search_duckduckgo(query, 4))  # Fewer from DDG as backup
        
        # Wait for both searches to complete with overall timeout
        try:
            brave_results, duckduckgo_results = await asyncio.wait_for(
                asyncio.gather(brave_task, duckduckgo_task, return_exceptions=True),
                timeout=12.0  # Overall timeout for both searches
            )
        except asyncio.TimeoutError:
            logger.warning("Combined search timed out, cancelling remaining tasks")
            brave_task.cancel()
            duckduckgo_task.cancel()
            brave_results, duckduckgo_results = [], []
        
        # Handle exceptions and log results
        if isinstance(brave_results, Exception):
            logger.error(f"Brave search failed: {brave_results}")
            brave_results = []
        else:
            logger.info(f"Brave search returned {len(brave_results)} results")
            
        if isinstance(duckduckgo_results, Exception):
            logger.error(f"DuckDuckGo search failed: {duckduckgo_results}")
            duckduckgo_results = []
        else:
            logger.info(f"DuckDuckGo search returned {len(duckduckgo_results)} results")
        
        # If both searches fail, return empty with warning
        if not brave_results and not duckduckgo_results:
            logger.warning(f"All search engines failed for query: {query}")
            return []
        
        # Combine results
        combined_results = brave_results + duckduckgo_results
        
        # Remove duplicates based on URL
        seen_urls = set()
        unique_results = []
        
        for result in combined_results:
            url = result.get("href", "")
            if url and url not in seen_urls:
                seen_urls.add(url)
                unique_results.append(result)
        
        # Return top results up to max_results
        final_results = unique_results[:max_results]
        logger.info(f"Combined Search: Returning {len(final_results)} total unique results")
        
        # Log search engine performance
        brave_count = len([r for r in final_results if r.get('source') == 'Brave'])
        ddg_count = len([r for r in final_results if r.get('source') == 'DuckDuckGo'])
        logger.info(f"Search distribution: Brave={brave_count}, DuckDuckGo={ddg_count}")
        
        return final_results
        
    except Exception as e:
        logger.error(f"Combined search error: {e}")
        return []


def preprocess_text(text: str) -> str:
    """Use spaCy for fast text cleaning/normalization"""
    doc = nlp(text)
    # Lemmatize and remove stopwords
    return " ".join([
        token.lemma_ for token in doc
        if not token.is_stop and not token.is_punct
    ])[:2048]

def format_conversation_history(history: Optional[List[Dict[str, str]]], max_entries: int = 10) -> str:
    """Format conversation history for inclusion in AI prompt"""
    if not history:
        return ""
    
    try:
        formatted_entries = []
        # Take the most recent entries up to max_entries
        recent_history = history[-max_entries:] if len(history) > max_entries else history
        
        for entry in recent_history:
            # Handle both formats: {"role": "user/assistant", "content": "..."} 
            # and {"user": "...", "assistant": "..."}
            if "role" in entry and "content" in entry:
                role = entry["role"].title()  # User or Assistant
                content = entry["content"].strip()
                if content:
                    formatted_entries.append(f"{role}: {content}")
            elif "user" in entry and "assistant" in entry:
                user_msg = entry["user"].strip()
                assistant_msg = entry["assistant"].strip()
                if user_msg and assistant_msg:
                    formatted_entries.append(f"User: {user_msg}")
                    formatted_entries.append(f"Assistant: {assistant_msg}")
        
        if formatted_entries:
            return "\n".join(formatted_entries)
        return ""
        
    except Exception as e:
        logger.warning(f"Error formatting conversation history: {e}")
        return ""






    
@app.on_event("startup")
async def startup_event():
    """Perform startup tasks like NLP model loading"""
    logger.info("Starting NLP model loading...")
    # spaCy and NLTK data are loaded implicitly on first use or by spacy.load
    
    # Test analytics database connection
    try:
        from analytics.database import test_connection
        analytics_connected = await test_connection()
        if analytics_connected:
            logger.info("Analytics database connection successful")
        else:
            logger.warning("Analytics database connection failed - analytics disabled")
    except Exception as e:
        logger.warning(f"Analytics initialization failed: {e}")
    
    logger.info("Startup complete - NLP models ready")

@app.post("/chat", response_model=ChatResponse)
async def chat_endpoint(request: ChatRequest, 
                       response: Response,
                       user_agent: str = Header(None),
                       x_session_id: str = Header(None)):
    """Enhanced chat endpoint with dynamic model selection and combined search"""
    logger.info(f"Request: {request.prompt}")
    
    # Validate and extract user_id
    user_id = validate_user_id(request.user_id)
    
    # Analytics setup
    session_id = x_session_id
    message_id = None
    start_time = None
    
    try:
        # Import analytics (with fallback if not available)
        try:
            from analytics.collectors import create_session, get_session, track_message, PerformanceTimer
            analytics_available = True
        except ImportError:
            logger.warning("Analytics not available")
            analytics_available = False
            
            # Create fallback PerformanceTimer when analytics unavailable
            class PerformanceTimer:
                def __init__(self):
                    self.start_time = None
                    self.end_time = None
                
                def __enter__(self):
                    import time
                    self.start_time = time.time()
                    return self
                
                def __exit__(self, exc_type, exc_val, exc_tb):
                    import time
                    self.end_time = time.time()
                
                @property
                def duration_ms(self) -> int:
                    if self.start_time and self.end_time:
                        return int((self.end_time - self.start_time) * 1000)
                    return 0
        
        # Start performance timing
        with PerformanceTimer() as timer:
            # Handle session management
            if analytics_available:
                if not session_id:
                    # Create new session
                    session = await create_session(user_agent=user_agent, user_id=user_id)
                    session_id = session.session_id
                else:
                    # Get existing session or create new one if not found
                    session = await get_session(session_id)
                    if not session:
                        session = await create_session(user_agent=user_agent, user_id=user_id)
                        session_id = session.session_id
            
            search_results = []
            search_context = ""
            search_decision = None
            cache_info = None
            
            # Phase 3: Get universal cache for result caching
            search_cache = get_universal_cache()
            
            # Phase 3b: Context-Aware Request Flow Optimization
            logger.info(f"Use Search : {request.use_search}")
            has_history = has_meaningful_conversation_history(request.history)
            logger.info(f"Has meaningful conversation history: {has_history}")
            
            if request.use_search:
                # Check if search should be forced (overrides all optimizations)
                if request.force_search:
                    search_decision = {
                        "should_search": True,
                        "reason": "Search forced by user",
                        "confidence": 1.0,
                        "flow_type": "forced"
                    }
                    perform_search = True
                    logger.info(f"Using forced search flow")
                elif not has_history:
                    # First Message (No History): Cache-first approach
                    logger.info(f"Using cache-first flow (no conversation history)")
                    search_terms = extract_search_terms(request.prompt.lower(), nlp, rake)
                    logger.info(f"Extract search terms successful: {search_terms}")
                    
                    # Try to get from cache first (skip search decision for performance)
                    cached_entry = search_cache.get(search_terms, use_semantic_matching=True, similarity_threshold=0.7)
                    
                    if cached_entry:
                        # Cache hit! Use cached results, no search needed
                        search_results = cached_entry.results
                        search_context = format_search_context(search_results)
                        cache_info = {
                            "cache_hit": True,
                            "cached_query": cached_entry.search_query,
                            "cache_age_seconds": int(time.time() - cached_entry.timestamp),
                            "hit_count": cached_entry.hit_count,
                            "flow_type": "cache_first_hit",
                            "cache_type": "chromadb_vector"
                        }
                        search_decision = {
                            "should_search": False,
                            "reason": "Cache hit in cache-first flow",
                            "confidence": 1.0,
                            "flow_type": "cache_first_hit",
                            "cache_type": "chromadb_vector"
                        }
                        perform_search = False
                        logger.info(f"ChromaDB Cache-first HIT: Using cached results (age: {cache_info['cache_age_seconds']}s, hits: {cached_entry.hit_count})")
                    else:
                        # Cache miss - perform web search without search decision overhead
                        logger.info(f"ChromaDB Cache-first MISS: Performing web search")
                        search_query = " ".join(search_terms) or request.prompt
                        search_results = await search_web_combined(search_query, 10)
                        search_context = format_search_context(search_results)
                        
                        # Store results in cache if search was successful
                        if search_results:
                            search_cache.put(search_terms, search_query, search_results)
                            logger.info(f"Cached search results in universal cache for future use")
                        
                        cache_info = {
                            "cache_hit": False,
                            "stored_in_cache": len(search_results) > 0,
                            "flow_type": "cache_first_miss",
                            "cache_type": "chromadb_vector"
                        }
                        search_decision = {
                            "should_search": True,
                            "reason": "Cache miss in cache-first flow",
                            "confidence": 1.0,
                            "flow_type": "cache_first_miss",
                            "cache_type": "chromadb_vector"
                        }
                        perform_search = True
                else:
                    # Follow-up Messages (Has History): Search decision first
                    logger.info(f"Using search-decision-first flow (has conversation history)")
                    # Use hybrid intelligent search decision (Phase 2: AI-enhanced)
                    search_decision = await hybrid_search_decision(
                        request.prompt, 
                        request.history, 
                        request.search_decision_mode,
                        nlp,
                        model
                    )
                    search_decision["flow_type"] = "search_decision_first"
                    perform_search = search_decision["should_search"]
                    
                    if perform_search:
                        # Search needed - check cache before web search
                        search_terms = extract_search_terms(request.prompt.lower(), nlp, rake)
                        logger.info(f"Extract search terms successful: {search_terms}")
                        
                        # Try to get from cache first
                        cached_entry = search_cache.get(search_terms, use_semantic_matching=True, similarity_threshold=0.7)
                        
                        if cached_entry:
                            # Cache hit! Use cached results
                            search_results = cached_entry.results
                            search_context = format_search_context(search_results)
                            cache_info = {
                                "cache_hit": True,
                                    "cached_query": cached_entry.search_query,
                                "cache_age_seconds": int(time.time() - cached_entry.timestamp),
                                "hit_count": cached_entry.hit_count,
                                "flow_type": "search_decision_cache_hit",
                                "cache_type": "chromadb_vector"
                            }
                            logger.info(f"ChromaDB Search-decision flow Cache HIT: Using cached results (age: {cache_info['cache_age_seconds']}s, hits: {cached_entry.hit_count})")
                        else:
                            # Cache miss - perform web search
                            logger.info(f"ChromaDB Search-decision flow Cache MISS: Performing web search")
                            search_query = " ".join(search_terms) or request.prompt
                            search_results = await search_web_combined(search_query, 10)
                            search_context = format_search_context(search_results)
                            
                            # Store results in cache if search was successful
                            if search_results:
                                search_cache.put(search_terms, search_query, search_results)
                                logger.info(f"Cached search results in universal cache for future use")
                            
                            cache_info = {
                                "cache_hit": False,
                                    "stored_in_cache": len(search_results) > 0,
                                "flow_type": "search_decision_cache_miss",
                                "cache_type": "chromadb_vector"
                            }
                    else:
                        # Search not needed based on conversation context
                        logger.info(f"Search skipped: {search_decision['reason']}")
                        cache_info = {
                            "search_skipped": True,
                            "flow_type": "search_decision_skip"
                        }
                
                logger.info(f"Search Decision: {search_decision}")
                
                if perform_search:
                    logger.info(f"Search processing complete - Results: {len(search_results)}")
            else:
                logger.info(f"Search disabled by request")
                cache_info = {"search_disabled": True}

            logger.info(f"Search Context: {search_context}")
            
            # Format conversation history
            conversation_history = format_conversation_history(request.history, max_entries=10)
            logger.info(f"Conversation History: {len(request.history or [])} entries")
            
            # Create a prompt for the AI model with conversation history
            if conversation_history:
                prompt_template = f"""
                You are having a conversation with a user. Here is the conversation history:

                Conversation History:
                ---
                {conversation_history}
                ---
                
                Based on the following context from web pages I have read and the conversation history above, please answer the user's question.
                If the context does not contain the answer and you cannot answer based on the conversation history, say that you don't have enough information.

                Context:
                ---
                {search_context}
                ---

                Question: {request.prompt}
                
                Answer:
                """
            else:
                prompt_template = f"""
                Based on the following context from web pages I have read, please answer the user's question.
                If the context does not contain the answer, say that you don't have enough information.

                Context:
                ---
                {search_context}
                ---

                Question: {request.prompt}
                
                Answer:
                """

            response_text = await run_gemini_inference(prompt_template)

        # Track message analytics
        if analytics_available and session_id:
            message = await track_message(
                session_id=session_id,
                prompt_length=len(request.prompt),
                response_length=len(response_text),
                response_time_ms=timer.duration_ms,
                used_search=request.use_search,
                max_tokens=request.max_new_tokens,
                temperature=request.temperature,
                success=True,
                user_id=user_id
            )
            if message:
                message_id = message.message_id

        # Add session ID to response headers
        if session_id:
            response.headers["X-Session-ID"] = session_id

        # Prepare response
        chat_response = ChatResponse(
            response=response_text,
            search_results=search_results,
            search_decision=search_decision,
            cache_info=cache_info
        )

        return chat_response

    except Exception as e:
        # Track failed message
        if analytics_available and session_id:
            await track_message(
                session_id=session_id,
                prompt_length=len(request.prompt),
                response_length=0,
                response_time_ms=timer.duration_ms if 'timer' in locals() else 0,
                used_search=request.use_search,
                max_tokens=request.max_new_tokens,
                temperature=request.temperature,
                success=False,
                error_message=str(e),
                user_id=user_id
            )
        
        logger.error(f"Chat error: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/search")
async def search_endpoint(request: SearchRequest):
    """Search endpoint with combined search engines"""
    try:
        results = await search_web_combined(request.query, request.max_results)
        return {"results": results}
    except Exception as e:
        logger.error(f"Search endpoint error: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/analytics/stats")
async def analytics_stats():
    """Get basic analytics statistics"""
    try:
        from analytics.dashboard import get_basic_stats
        stats = await get_basic_stats()
        return stats
    except Exception as e:
        logger.error(f"Analytics stats error: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/analytics/dashboard")
async def analytics_dashboard():
    """Get HTML analytics dashboard"""
    try:
        from analytics.dashboard import get_dashboard_data
        from fastapi.responses import HTMLResponse
        
        # Get dashboard data including user metrics
        data = await get_dashboard_data()
        
        # Get user statistics for the dashboard
        from analytics.dashboard import get_user_statistics, get_authenticated_vs_anonymous_metrics
        user_stats = await get_user_statistics()
        comparison_stats = await get_authenticated_vs_anonymous_metrics()
        
        # Create HTML dashboard
        html_content = f"""
        <!DOCTYPE html>
        <html lang="en">
        <head>
            <meta charset="UTF-8">
            <meta name="viewport" content="width=device-width, initial-scale=1.0">
            <title>Atlas Analytics Dashboard</title>
            <script src="https://cdn.jsdelivr.net/npm/chart.js"></script>
            <style>
                body {{
                    font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
                    margin: 0;
                    padding: 20px;
                    background-color: #f5f5f5;
                }}
                .container {{
                    max-width: 1200px;
                    margin: 0 auto;
                }}
                .header {{
                    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
                    color: white;
                    padding: 30px;
                    border-radius: 10px;
                    margin-bottom: 30px;
                    text-align: center;
                }}
                .stats-grid {{
                    display: grid;
                    grid-template-columns: repeat(auto-fit, minmax(250px, 1fr));
                    gap: 20px;
                    margin-bottom: 30px;
                }}
                .stat-card {{
                    background: white;
                    padding: 25px;
                    border-radius: 10px;
                    box-shadow: 0 2px 10px rgba(0,0,0,0.1);
                    text-align: center;
                }}
                .stat-number {{
                    font-size: 2.5em;
                    font-weight: bold;
                    color: #667eea;
                    margin-bottom: 10px;
                }}
                .stat-label {{
                    color: #666;
                    font-size: 0.9em;
                    text-transform: uppercase;
                    letter-spacing: 1px;
                }}
                .chart-container {{
                    background: white;
                    padding: 25px;
                    border-radius: 10px;
                    box-shadow: 0 2px 10px rgba(0,0,0,0.1);
                    margin-bottom: 20px;
                }}
                .chart-title {{
                    font-size: 1.2em;
                    font-weight: bold;
                    margin-bottom: 20px;
                    color: #333;
                }}
                .refresh-btn {{
                    background: #667eea;
                    color: white;
                    border: none;
                    padding: 10px 20px;
                    border-radius: 5px;
                    cursor: pointer;
                    font-size: 1em;
                    margin-bottom: 20px;
                }}
                .refresh-btn:hover {{
                    background: #5a6fd8;
                }}
                .error {{
                    background: #fee;
                    color: #c33;
                    padding: 15px;
                    border-radius: 5px;
                    margin: 10px 0;
                }}
                .last-updated {{
                    text-align: center;
                    color: #666;
                    font-size: 0.9em;
                    margin-top: 20px;
                }}
            </style>
        </head>
        <body>
            <div class="container">
                <div class="header">
                    <h1>🚀 Atlas Analytics Dashboard</h1>
                    <p>Real-time insights into your chat application</p>
                </div>
                
                <button class="refresh-btn" onclick="location.reload()">🔄 Refresh Data</button>
                
                <div class="stats-grid">
                    <div class="stat-card">
                        <div class="stat-number">{data.get('basic', {}).get('total_messages', 0)}</div>
                        <div class="stat-label">Total Messages</div>
                    </div>
                    <div class="stat-card">
                        <div class="stat-number">{data.get('basic', {}).get('total_sessions', 0)}</div>
                        <div class="stat-label">Total Sessions</div>
                    </div>
                    <div class="stat-card">
                        <div class="stat-number">{data.get('basic', {}).get('active_sessions', 0)}</div>
                        <div class="stat-label">Active Sessions</div>
                    </div>
                    <div class="stat-card">
                        <div class="stat-number">{data.get('basic', {}).get('messages_today', 0)}</div>
                        <div class="stat-label">Messages Today</div>
                    </div>
                    <div class="stat-card">
                        <div class="stat-number">{data.get('basic', {}).get('search_usage_percentage', 0)}%</div>
                        <div class="stat-label">Search Usage</div>
                    </div>
                    <div class="stat-card">
                        <div class="stat-number">{data.get('basic', {}).get('average_response_time_ms', 0)}ms</div>
                        <div class="stat-label">Avg Response Time</div>
                    </div>
                </div>
                
                <div class="chart-container">
                    <div class="chart-title">🔓 Anonymous Usage Overview</div>
                    <div class="stats-grid">
                        <div class="stat-card">
                            <div class="stat-number">{user_stats.get('anonymous_sessions', 0)}</div>
                            <div class="stat-label">Anonymous Sessions</div>
                        </div>
                        <div class="stat-card">
                            <div class="stat-number">{user_stats.get('anonymous_messages', 0)}</div>
                            <div class="stat-label">Anonymous Messages</div>
                        </div>
                        <div class="stat-card">
                            <div class="stat-number">{round(100 - user_stats.get('authenticated_session_percentage', 0), 1)}%</div>
                            <div class="stat-label">Anonymous Session %</div>
                        </div>
                        <div class="stat-card">
                            <div class="stat-number">{round(100 - user_stats.get('authenticated_message_percentage', 0), 1)}%</div>
                            <div class="stat-label">Anonymous Message %</div>
                        </div>
                    </div>
                </div>
                
                <div class="chart-container">
                    <div class="chart-title">👥 User Analytics</div>
                    <div class="stats-grid">
                        <div class="stat-card">
                            <div class="stat-number">{user_stats.get('unique_authenticated_users', 0)}</div>
                            <div class="stat-label">Unique Users</div>
                        </div>
                        <div class="stat-card">
                            <div class="stat-number">{user_stats.get('authenticated_sessions', 0)}</div>
                            <div class="stat-label">Authenticated Sessions</div>
                        </div>
                        <div class="stat-card">
                            <div class="stat-number">{user_stats.get('anonymous_sessions', 0)}</div>
                            <div class="stat-label">Anonymous Sessions</div>
                        </div>
                        <div class="stat-card">
                            <div class="stat-number">{user_stats.get('authenticated_session_percentage', 0)}%</div>
                            <div class="stat-label">Auth Session %</div>
                        </div>
                        <div class="stat-card">
                            <div class="stat-number">{user_stats.get('authenticated_messages', 0)}</div>
                            <div class="stat-label">Authenticated Messages</div>
                        </div>
                        <div class="stat-card">
                            <div class="stat-number">{user_stats.get('anonymous_messages', 0)}</div>
                            <div class="stat-label">Anonymous Messages</div>
                        </div>
                        <div class="stat-card">
                            <div class="stat-number">{user_stats.get('authenticated_message_percentage', 0)}%</div>
                            <div class="stat-label">Auth Message %</div>
                        </div>
                        <div class="stat-card">
                            <div class="stat-number">{user_stats.get('anonymous_messages', 0) + user_stats.get('authenticated_messages', 0)}</div>
                            <div class="stat-label">Total Messages</div>
                        </div>
                    </div>
                </div>
                
                <div class="chart-container">
                    <div class="chart-title">🔍 User Filtering</div>
                    <div style="margin-bottom: 20px;">
                        <input type="text" id="userIdInput" placeholder="Enter user ID to filter analytics" 
                               style="padding: 10px; border: 1px solid #ddd; border-radius: 5px; width: 300px; margin-right: 10px;">
                        <button onclick="filterByUser()" style="padding: 10px 20px; background: #667eea; color: white; border: none; border-radius: 5px; cursor: pointer;">
                            Filter Analytics
                        </button>
                        <button onclick="clearFilter()" style="padding: 10px 20px; background: #6c757d; color: white; border: none; border-radius: 5px; cursor: pointer; margin-left: 10px;">
                            Clear Filter
                        </button>
                    </div>
                    <div id="userFilterResults" style="display: none;">
                        <h4>User-Specific Analytics</h4>
                        <div id="userStatsGrid" class="stats-grid"></div>
                    </div>
                </div>
                
                <div class="chart-container">
                    <div class="chart-title">📊 Hourly Message Activity (Last 24 Hours)</div>
                    <canvas id="hourlyChart" width="400" height="200"></canvas>
                </div>
                
                <div class="chart-container">
                    <div class="chart-title">⚡ Performance Metrics</div>
                    <canvas id="performanceChart" width="400" height="200"></canvas>
                </div>
                
                <div class="chart-container">
                    <div class="chart-title">👤 Authenticated vs Anonymous Comparison</div>
                    <canvas id="comparisonChart" width="400" height="200"></canvas>
                </div>
                
                <div class="last-updated">
                    Last updated: {data.get('generated_at', 'Unknown')}
                </div>
            </div>
            
            <script>
                // Hourly Chart
                const hourlyData = {data.get('hourly', [])};
                const hourlyLabels = hourlyData.map(d => d.hour);
                const hourlyMessages = hourlyData.map(d => d.message_count);
                const hourlySearches = hourlyData.map(d => d.search_count);
                
                new Chart(document.getElementById('hourlyChart'), {{
                    type: 'line',
                    data: {{
                        labels: hourlyLabels,
                        datasets: [{{
                            label: 'Messages',
                            data: hourlyMessages,
                            borderColor: '#667eea',
                            backgroundColor: 'rgba(102, 126, 234, 0.1)',
                            tension: 0.4
                        }}, {{
                            label: 'With Search',
                            data: hourlySearches,
                            borderColor: '#f093fb',
                            backgroundColor: 'rgba(240, 147, 251, 0.1)',
                            tension: 0.4
                        }}]
                    }},
                    options: {{
                        responsive: true,
                        scales: {{
                            y: {{
                                beginAtZero: true
                            }}
                        }}
                    }}
                }});
                
                // Performance Chart
                const perfData = {data.get('performance', {})};
                new Chart(document.getElementById('performanceChart'), {{
                    type: 'bar',
                    data: {{
                        labels: ['P50', 'P90', 'P95', 'Error Rate %'],
                        datasets: [{{
                            label: 'Performance Metrics',
                            data: [
                                perfData.response_time_p50 || 0,
                                perfData.response_time_p90 || 0,
                                perfData.response_time_p95 || 0,
                                perfData.error_rate_percentage || 0
                            ],
                            backgroundColor: [
                                'rgba(102, 126, 234, 0.8)',
                                'rgba(240, 147, 251, 0.8)',
                                'rgba(255, 159, 64, 0.8)',
                                'rgba(255, 99, 132, 0.8)'
                            ]
                        }}]
                    }},
                    options: {{
                        responsive: true,
                        scales: {{
                            y: {{
                                beginAtZero: true
                            }}
                        }}
                    }}
                }});
                
                // Comparison Chart
                const comparisonData = {comparison_stats};
                new Chart(document.getElementById('comparisonChart'), {{
                    type: 'bar',
                    data: {{
                        labels: ['Sessions', 'Messages', 'Avg Response Time (ms)', 'Search Usage %'],
                        datasets: [{{
                            label: 'Authenticated Users',
                            data: [
                                comparisonData.authenticated?.sessions || 0,
                                comparisonData.authenticated?.messages || 0,
                                comparisonData.authenticated?.avg_response_time_ms || 0,
                                comparisonData.authenticated?.search_usage_percentage || 0
                            ],
                            backgroundColor: 'rgba(102, 126, 234, 0.8)'
                        }}, {{
                            label: 'Anonymous Users',
                            data: [
                                comparisonData.anonymous?.sessions || 0,
                                comparisonData.anonymous?.messages || 0,
                                comparisonData.anonymous?.avg_response_time_ms || 0,
                                comparisonData.anonymous?.search_usage_percentage || 0
                            ],
                            backgroundColor: 'rgba(255, 159, 64, 0.8)'
                        }}]
                    }},
                    options: {{
                        responsive: true,
                        scales: {{
                            y: {{
                                beginAtZero: true
                            }}
                        }}
                    }}
                }});
                
                // User filtering functions
                async function filterByUser() {{
                    const userId = document.getElementById('userIdInput').value.trim();
                    if (!userId) {{
                        alert('Please enter a user ID');
                        return;
                    }}
                    
                    try {{
                        const response = await fetch(`/analytics/user/${{encodeURIComponent(userId)}}`);
                        const userData = await response.json();
                        
                        if (userData.error) {{
                            alert(`Error: ${{userData.error}}`);
                            return;
                        }}
                        
                        displayUserStats(userData);
                    }} catch (error) {{
                        alert(`Error fetching user data: ${{error.message}}`);
                    }}
                }}
                
                function displayUserStats(userData) {{
                    const resultsDiv = document.getElementById('userFilterResults');
                    const statsGrid = document.getElementById('userStatsGrid');
                    
                    statsGrid.innerHTML = `
                        <div class="stat-card">
                            <div class="stat-number">${{userData.total_sessions || 0}}</div>
                            <div class="stat-label">User Sessions</div>
                        </div>
                        <div class="stat-card">
                            <div class="stat-number">${{userData.total_messages || 0}}</div>
                            <div class="stat-label">User Messages</div>
                        </div>
                        <div class="stat-card">
                            <div class="stat-number">${{userData.search_usage_percentage || 0}}%</div>
                            <div class="stat-label">Search Usage</div>
                        </div>
                        <div class="stat-card">
                            <div class="stat-number">${{userData.avg_response_time_ms || 0}}ms</div>
                            <div class="stat-label">Avg Response Time</div>
                        </div>
                        <div class="stat-card">
                            <div class="stat-number">${{userData.avg_messages_per_session || 0}}</div>
                            <div class="stat-label">Avg Msgs/Session</div>
                        </div>
                        <div class="stat-card">
                            <div class="stat-number">${{userData.active_sessions || 0}}</div>
                            <div class="stat-label">Active Sessions</div>
                        </div>
                    `;
                    
                    resultsDiv.style.display = 'block';
                }}
                
                function clearFilter() {{
                    document.getElementById('userIdInput').value = '';
                    document.getElementById('userFilterResults').style.display = 'none';
                }}
            </script>
        </body>
        </html>
        """
        
        return HTMLResponse(content=html_content)
        
    except Exception as e:
        logger.error(f"Analytics dashboard error: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/analytics/users")
async def analytics_users():
    """Get overall user statistics including authenticated vs anonymous metrics"""
    try:
        from analytics.dashboard import get_user_statistics
        stats = await get_user_statistics()
        return stats
    except Exception as e:
        logger.error(f"Analytics users error: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/analytics/user/{user_id}")
async def analytics_user(user_id: str):
    """Get analytics for a specific user"""
    try:
        # Validate user_id
        if not user_id or not isinstance(user_id, str) or len(user_id.strip()) == 0:
            raise HTTPException(status_code=400, detail="Invalid user_id provided")
        
        from analytics.dashboard import get_user_analytics
        stats = await get_user_analytics(user_id.strip())
        return stats
    except Exception as e:
        logger.error(f"Analytics user error: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/analytics/comparison")
async def analytics_comparison():
    """Get detailed comparison metrics between authenticated and anonymous users"""
    try:
        from analytics.dashboard import get_authenticated_vs_anonymous_metrics
        stats = await get_authenticated_vs_anonymous_metrics()
        return stats
    except Exception as e:
        logger.error(f"Analytics comparison error: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/analytics/export")
async def analytics_export(format: str = "json", days: int = 7, user_id: Optional[str] = None):
    """Export analytics data in JSON or CSV format with optional user_id filtering"""
    try:
        from analytics.database import get_sessions_collection, get_messages_collection
        from fastapi.responses import StreamingResponse
        from datetime import datetime, timedelta
        import json
        import csv
        import io
        
        # Validate format
        if format not in ["json", "csv"]:
            raise HTTPException(status_code=400, detail="Format must be 'json' or 'csv'")
        
        # Calculate date range
        end_date = datetime.utcnow()
        start_date = end_date - timedelta(days=days)
        
        # Get collections
        sessions_collection = await get_sessions_collection()
        messages_collection = await get_messages_collection()
        
        if sessions_collection is None or messages_collection is None:
            raise HTTPException(status_code=500, detail="Database not available")
        
        # Build filters with optional user_id
        session_filter = {"start_time": {"$gte": start_date, "$lte": end_date}}
        message_filter = {"timestamp": {"$gte": start_date, "$lte": end_date}}
        
        if user_id is not None and user_id.strip():
            user_id = user_id.strip()
            session_filter["user_id"] = user_id
            message_filter["user_id"] = user_id
        
        # Get data
        sessions_cursor = sessions_collection.find(session_filter)
        sessions_data = await sessions_cursor.to_list(None)
        
        messages_cursor = messages_collection.find(message_filter)
        messages_data = await messages_cursor.to_list(None)
        
        # Convert ObjectId to string for JSON serialization
        for session in sessions_data:
            session["_id"] = str(session["_id"])
            if "start_time" in session:
                session["start_time"] = session["start_time"].isoformat()
            if "end_time" in session and session["end_time"]:
                session["end_time"] = session["end_time"].isoformat()
        
        for message in messages_data:
            message["_id"] = str(message["_id"])
            if "timestamp" in message:
                message["timestamp"] = message["timestamp"].isoformat()
        
        export_data = {
            "export_info": {
                "generated_at": end_date.isoformat(),
                "date_range": {
                    "start": start_date.isoformat(),
                    "end": end_date.isoformat(),
                    "days": days
                },
                "filters": {
                    "user_id": user_id if user_id and user_id.strip() else None
                },
                "counts": {
                    "sessions": len(sessions_data),
                    "messages": len(messages_data)
                }
            },
            "sessions": sessions_data,
            "messages": messages_data
        }
        
        if format == "json":
            # Return JSON
            json_str = json.dumps(export_data, indent=2, default=str)
            
            def generate():
                yield json_str
            
            # Generate filename with optional user_id
            filename_suffix = f"_user_{user_id}" if user_id and user_id.strip() else ""
            filename = f"atlas_analytics_{start_date.strftime('%Y%m%d')}_{end_date.strftime('%Y%m%d')}{filename_suffix}.json"
            
            return StreamingResponse(
                generate(),
                media_type="application/json",
                headers={"Content-Disposition": f"attachment; filename={filename}"}
            )
        
        elif format == "csv":
            # Create CSV with separate sheets for sessions and messages
            output = io.StringIO()
            
            # Write export info
            output.write(f"# Atlas Analytics Export\\n")
            output.write(f"# Generated: {export_data['export_info']['generated_at']}\\n")
            output.write(f"# Date Range: {export_data['export_info']['date_range']['start']} to {export_data['export_info']['date_range']['end']}\\n")
            output.write(f"# Sessions: {export_data['export_info']['counts']['sessions']}, Messages: {export_data['export_info']['counts']['messages']}\\n")
            output.write("\\n")
            
            # Sessions CSV
            output.write("=== SESSIONS ===\\n")
            if sessions_data:
                fieldnames = sessions_data[0].keys()
                writer = csv.DictWriter(output, fieldnames=fieldnames)
                writer.writeheader()
                writer.writerows(sessions_data)
            
            output.write("\\n=== MESSAGES ===\\n")
            if messages_data:
                fieldnames = messages_data[0].keys()
                writer = csv.DictWriter(output, fieldnames=fieldnames)
                writer.writeheader()
                writer.writerows(messages_data)
            
            def generate():
                yield output.getvalue()
            
            # Generate filename with optional user_id
            filename_suffix = f"_user_{user_id}" if user_id and user_id.strip() else ""
            filename = f"atlas_analytics_{start_date.strftime('%Y%m%d')}_{end_date.strftime('%Y%m%d')}{filename_suffix}.csv"
            
            return StreamingResponse(
                generate(),
                media_type="text/csv",
                headers={"Content-Disposition": f"attachment; filename={filename}"}
            )
        
    except Exception as e:
        logger.error(f"Analytics export error: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/analytics/cache")
async def analytics_cache():
    """Get universal ChromaDB cache analytics and performance metrics"""
    try:
        # Get universal cache instance
        search_cache = get_universal_cache()
        
        # Get cache stats
        cache_stats = search_cache.get_stats()
        
        # Get popular queries from ChromaDB
        popular_queries = search_cache.get_popular_queries(limit=10)
        
        # Determine cache effectiveness
        hit_rate = cache_stats.get("hit_rate_percentage", 0)
        cache_effectiveness = "High" if hit_rate > 60 else "Medium" if hit_rate > 30 else "Low"
        
        return {
            "cache_statistics": cache_stats,
            "popular_queries": popular_queries,
            "cache_effectiveness": cache_effectiveness,
            "memory_efficiency": {
                "entries_per_mb": cache_stats["cache_size"] / max(0.1, cache_stats["memory_usage_mb"]),
                "avg_entry_size_kb": (cache_stats["memory_usage_mb"] * 1024) / max(1, cache_stats["cache_size"])
            },
            "vector_database_info": {
                "embedding_model": cache_stats.get("embedding_model", "unknown"),
                "similarity_threshold": cache_stats.get("similarity_threshold", 0.7),
                "persistent_storage": cache_stats.get("persistent_storage", False),
                "database_path": cache_stats.get("database_path", "unknown")
            }
        }
            
    except Exception as e:
        logger.error(f"ChromaDB cache analytics error: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/analytics/cache/clear")
async def clear_cache(cache_type: str = "expired"):
    """Clear universal ChromaDB cache data - for maintenance and testing"""
    try:
        if cache_type not in ["expired", "all"]:
            raise HTTPException(status_code=400, detail="cache_type must be 'expired' or 'all'")
        
        # Get cache instance
        search_cache = get_universal_cache()
        
        cache_stats_before = search_cache.get_stats()
        
        if cache_type == "expired":
            search_cache.clear_expired()
            action_taken = "Cleared expired entries from ChromaDB universal cache"
        elif cache_type == "all":
            search_cache.clear_all()
            action_taken = "Cleared all entries from ChromaDB universal cache"
        
        cache_stats_after = search_cache.get_stats()
        
        return {
            "action": cache_type,
            "message": action_taken,
            "cache_type": "chromadb_vector",
            "before": {
                "cache_size": cache_stats_before["cache_size"],
                "memory_usage_mb": cache_stats_before["memory_usage_mb"]
            },
            "after": {
                "cache_size": cache_stats_after["cache_size"],
                "memory_usage_mb": cache_stats_after["memory_usage_mb"]
            },
            "entries_removed": cache_stats_before["cache_size"] - cache_stats_after["cache_size"],
            "memory_freed_mb": cache_stats_before["memory_usage_mb"] - cache_stats_after["memory_usage_mb"]
        }
        
    except Exception as e:
        logger.error(f"ChromaDB cache clear error: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/")
async def root():
    """Enhanced health check with model status"""
    return {
        "message": "Enhanced Chat API with Gemini and Combined Search is running!",
        "models": {
            "gemini_loaded": True, # Gemini is always loaded via API
            "model_type": "gemini-1.5-flash",
            "capabilities": ["question_answering", "summarization", "explanation", "web_search_augmentation"]
        },
        "search_engines": ["Brave", "DuckDuckGo"],
        "endpoints": {
            "chat": "/chat",
            "search": "/search",
            "analytics_stats": "/analytics/stats",
            "analytics_dashboard": "/analytics/dashboard",
            "analytics_users": "/analytics/users",
            "analytics_user": "/analytics/user/{user_id}",
            "analytics_comparison": "/analytics/comparison",
            "analytics_export": "/analytics/export",
            "analytics_cache": "/analytics/cache",
            "cache_clear": "/analytics/cache/clear",
            "docs": "/docs"
        },
        "cache_system": {
            "enabled": True,
            "cache_type": "chromadb_vector",
            "cache_size": get_universal_cache().get_stats().get("cache_size", 0),
            "max_size": get_universal_cache().max_size,
            "semantic_similarity": True,
            "vector_similarity": True,
            "persistent_storage": True,
            "ttl_management": True,
            "memory_efficient": True,
            "embedding_model": get_universal_cache().get_stats().get("embedding_model", "all-MiniLM-L6-v2")
        }
    }


async def run_gemini_inference(prompt_text: str) -> str:
    """Run Gemini model inference"""
    try:
        response = await model.generate_content_async(prompt_text)
        return response.text
    except Exception as e:
        logger.error(f"Gemini inference error: {e}")
        raise


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)