app.py

import os
import logging
import time
from datetime import datetime
from contextlib import asynccontextmanager
from fastapi import FastAPI, WebSocket, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, FileResponse
from websocket_handler import handle_websocket_connection
from enhanced_websocket_handler import handle_enhanced_websocket_connection
from conversational_websocket_handler import handle_conversational_websocket
from hybrid_llm_service import HybridLLMService
from voice_service import VoiceService
from groq_voice_service import groq_voice_service  # Import the new Groq voice service
from rag_service import search_documents_async
from lancedb_service import LanceDBService
from scenario_analysis_service import ScenarioAnalysisService
from evidence_pack_export import export_evidence_pack_pdf, export_evidence_pack_csv
import config
from dotenv import load_dotenv

# MCP and Authentication imports
from fastapi import Depends
from pydantic import BaseModel
from typing import Optional
from auth import get_current_user

# Load environment variables
load_dotenv()

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger(__name__)

# Get configuration
config_dict = {
    "ALLOWED_ORIGINS": config.ALLOWED_ORIGINS,
    "ENABLE_VOICE_FEATURES": config.ENABLE_VOICE_FEATURES
}

@asynccontextmanager
async def lifespan(app: FastAPI):
    """Application lifespan handler"""
    # Startup
    logger.info("🚀 Starting Voice Bot Application...")
    
    # Setup sample documents if database is empty
    try:
        from setup_documents import setup_sample_documents
        await setup_sample_documents()
    except Exception as e:
        logger.warning(f"⚠️ Could not setup sample documents: {e}")
    
    logger.info("✅ Application started successfully")
    yield
    # Shutdown (if needed)
    logger.info("🛑 Shutting down Voice Bot Application...")

# Create FastAPI application
app = FastAPI(
    title="Voice Bot Government Assistant",
    description="AI-powered voice assistant for government policies and services",
    version="1.0.0",
    lifespan=lifespan
)

# Configure CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=config.ALLOWED_ORIGINS,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Initialize services (lazy loading for HF Spaces)
llm_service = None
voice_service = None
lancedb_service = None
scenario_service = None

def get_llm_service():
    global llm_service
    if llm_service is None:
        llm_service = HybridLLMService()
    return llm_service

def get_voice_service():
    global voice_service
    if voice_service is None:
        voice_service = VoiceService()
    return voice_service

def get_lancedb_service():
    global lancedb_service
    if lancedb_service is None:
        lancedb_service = LanceDBService()
    return lancedb_service

def get_scenario_service():
    global scenario_service
    if scenario_service is None:
        scenario_service = ScenarioAnalysisService()
    return scenario_service

# Evidence Pack Export endpoint
@app.api_route("/export_evidence_pack", methods=["GET", "POST", "OPTIONS", "HEAD"])
async def export_evidence_pack(request: Request, format: str = "pdf"):
    """Export evidence pack in PDF or CSV format"""
    try:
        # Handle CORS preflight and HEAD requests
        if request.method == "OPTIONS":
            return JSONResponse({"status": "ok"}, status_code=200)
        
        if request.method == "HEAD":
            # For HEAD requests, return headers without body
            return JSONResponse(
                {"status": "ok"}, 
                status_code=200,
                headers={"Content-Type": "application/pdf" if format.lower() == "pdf" else "text/csv"}
            )
        
        # Handle both GET and POST requests
        if request.method == "POST":
            try:
                data = await request.json()
                # Format can come from query params or request body
                format = request.query_params.get("format", data.get("format", "pdf"))
            except Exception:
                # If JSON parsing fails, use query params
                data = {
                    "query": request.query_params.get("query", ""),
                    "format": format,
                    "timestamp": datetime.now().isoformat()
                }
        else:  # GET request
            # For GET requests, we need some default data structure
            data = {
                "query": request.query_params.get("query", ""),
                "format": format,
                "timestamp": datetime.now().isoformat(),
                "message": "Sample evidence pack export"
            }
        
        if format.lower() == "pdf":
            file_path = export_evidence_pack_pdf(data)
            return FileResponse(
                file_path, 
                media_type="application/pdf", 
                filename="evidence_pack.pdf",
                headers={"Content-Disposition": "attachment; filename=evidence_pack.pdf"}
            )
        elif format.lower() == "csv":
            file_path = export_evidence_pack_csv(data)
            return FileResponse(
                file_path, 
                media_type="text/csv", 
                filename="evidence_pack.csv",
                headers={"Content-Disposition": "attachment; filename=evidence_pack.csv"}
            )
        else:
            return JSONResponse({"error": "Invalid format. Use 'pdf' or 'csv'"}, status_code=400)
    except Exception as e:
        logger.error(f"Export evidence pack error: {str(e)}")
        return JSONResponse({"error": f"Export failed: {str(e)}"}, status_code=500)

# Health check endpoint
@app.get("/health")
async def health_check():
    """Health check endpoint"""
    return {
        "status": "healthy",
        "service": "voice-bot-api",
        "timestamp": datetime.now().isoformat(),
        "version": "1.0.0"
    }

# Root endpoint
@app.get("/")
async def root():
    """Root endpoint with service information"""
    return {
        "message": "Voice Bot Government Assistant API",
        "status": "running",
        "version": "1.0.0",
        "endpoints": {
            "health": "/health",
            "chat": "/chat",
            "websocket": "/ws",
            "websocket_stream": "/ws/stream",
            "export_evidence_pack": "/export_evidence_pack",
            "docs": "/docs"
        }
    }

# Chat endpoint
@app.post("/chat")
async def chat_endpoint(request: dict):
    """Text-based chat endpoint"""
    try:
        message = request.get("message", "")
        if not message:
            raise HTTPException(status_code=400, detail="Message is required")
        
        llm = get_llm_service()
        response = await llm.get_response(message)
        
        return {
            "response": response,
            "timestamp": datetime.now().isoformat()
        }
    except Exception as e:
        logger.error(f"Chat error: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

# WebSocket endpoints
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    """WebSocket endpoint for real-time communication"""
    await handle_enhanced_websocket_connection(websocket)

@app.websocket("/ws/stream")
async def websocket_stream_endpoint(websocket: WebSocket):
    """
    Enhanced WebSocket endpoint with Groq ASR for superior voice transcription
    Based on friend's superior implementation for better accuracy
    """
    from groq_websocket_handler import groq_websocket_handler
    import json
    
    # Accept connection and get session ID
    session_id = await groq_websocket_handler.connect(websocket)
    
    try:
        while True:
            # Receive message from client
            message_text = await websocket.receive_text()
            
            try:
                message = json.loads(message_text)
            except json.JSONDecodeError:
                await groq_websocket_handler.send_message(session_id, {
                    "type": "error",
                    "message": "Invalid JSON message",
                    "timestamp": datetime.now().isoformat()
                })
                continue
            
            # Handle different message types
            await groq_websocket_handler.handle_stream_message(websocket, session_id, message)
            
    except Exception as e:
        logger.error(f"❌ WebSocket stream error: {e}")
    finally:
        await groq_websocket_handler.disconnect(session_id)

@app.websocket("/ws/conversational")
async def websocket_conversational_endpoint(websocket: WebSocket):
    """
    Enhanced Conversational WebSocket endpoint with session memory and personalization
    Based on friend's conversational implementation for better user experience
    """
    await handle_conversational_websocket(websocket)

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