app.py

import os
import json
import logging
import asyncio
import sys
from datetime import datetime, timedelta
from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect, BackgroundTasks, Request, Query
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from pydantic import BaseModel
from dotenv import load_dotenv
import time
from functools import lru_cache
from contextlib import asynccontextmanager
from typing import Dict, List, Optional, Any, Protocol, Callable, Union
import aiohttp
import requests
from db_connector import save_session, get_unanswered_questions, get_unanswered_question_by_session_id, find_user_question, collection

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# Load environment variables
load_dotenv()

# Notification service abstraction
class NotificationService:
    """
    Abstract interface for sending notifications. Allows using different notification methods.
    """
    async def send_notification(self, message: str) -> bool:
        """
        Send notification with message content.
        Returns True if successful, False if failed.
        """
        raise NotImplementedError("Subclasses must implement send_notification")

# Telegram Notification Service - used if admin_bot is available
class TelegramNotificationService(NotificationService):
    def __init__(self):
        # Instead of importing from admin_bot, we define the notification function ourselves
        self.bot_token = os.getenv("ADMIN_TELEGRAM_BOT_TOKEN")
        self.admin_chat_id = os.getenv("ADMIN_GROUP_CHAT_ID")
        
        if self.bot_token and self.admin_chat_id:
            logger.info("TelegramNotificationService: Telegram credentials configured")
            self.available = True
        else:
            logger.warning("TelegramNotificationService: No telegram bot token or admin chat id configured")
            self.available = False
    
    async def send_notification(self, message: str) -> bool:
        if not self.available:
            logger.warning("TelegramNotificationService not available")
            return False
        
        try:
            # Define function to send notification via Telegram API
            url = f"https://api.telegram.org/bot{self.bot_token}/sendMessage"
            payload = {
                "chat_id": self.admin_chat_id,
                "text": message,
                "parse_mode": "Markdown"
            }
            
            async with aiohttp.ClientSession() as session:
                async with session.post(url, json=payload) as response:
                    if response.status == 200:
                        result = await response.json()
                        logger.info(f"Notification sent via Telegram: message_id={result.get('result', {}).get('message_id')}")
                        return True
                    else:
                        error_data = await response.text()
                        logger.error(f"Error sending notification via Telegram API: {response.status}, {error_data}")
                        return False
        except Exception as e:
            logger.error(f"Error sending notification via Telegram: {e}")
            return False

# Webhook Notification Service - can be used when API is deployed separately
class WebhookNotificationService(NotificationService):
    def __init__(self, webhook_url: str = None):
        self.webhook_url = webhook_url or os.getenv("NOTIFICATION_WEBHOOK_URL")
        if self.webhook_url:
            logger.info(f"WebhookNotificationService: Webhook URL configured: {self.webhook_url}")
            self.available = True
        else:
            logger.warning("WebhookNotificationService: No webhook URL configured")
            self.available = False
    
    async def send_notification(self, message: str) -> bool:
        if not self.available:
            logger.warning("WebhookNotificationService not available")
            return False
        
        try:
            async with aiohttp.ClientSession() as session:
                async with session.post(self.webhook_url, json={"message": message}) as response:
                    if 200 <= response.status < 300:
                        logger.info(f"Notification sent to webhook: {response.status}")
                        return True
                    else:
                        logger.error(f"Error sending notification to webhook: {response.status}")
                        return False
        except Exception as e:
            logger.error(f"Error sending notification via webhook: {e}")
            return False

# Logging Notification Service - fallback method that just logs
class LoggingNotificationService(NotificationService):
    async def send_notification(self, message: str) -> bool:
        logger.info(f"NOTIFICATION: {message}")
        return True

# Cache Notification Service - store notifications for later retrieval
class CacheNotificationService(NotificationService):
    def __init__(self):
        self.notifications = []
    
    async def send_notification(self, message: str) -> bool:
        self.notifications.append({
            "message": message,
            "timestamp": datetime.now().isoformat()
        })
        logger.info(f"Notification saved to cache, currently {len(self.notifications)} notifications")
        return True
    
    def get_notifications(self, limit: int = 100):
        """Get cached notifications"""
        return self.notifications[-limit:]

# Service locator pattern
class NotificationManager:
    def __init__(self):
        self.services = []
        
        # Create notification services with priority
        # Priority order: Telegram > Webhook > Cache > Logging
        # Initialize notification services
        webhook_url = os.getenv("NOTIFICATION_WEBHOOK_URL")
        
        # 1. Telegram service (if available)
        telegram_service = TelegramNotificationService()
        if telegram_service.available:
            self.services.append(telegram_service)
        
        # 2. Webhook service (if configured)
        if webhook_url:
            webhook_service = WebhookNotificationService(webhook_url)
            self.services.append(webhook_service)
        
        # 3. Cache service (always available)
        self.cache_service = CacheNotificationService()
        self.services.append(self.cache_service)
        
        # 4. Logging service (always available)
        self.services.append(LoggingNotificationService())
        
        logger.info(f"Initialized {len(self.services)} notification services")
    
    async def send_notification(self, message: str) -> bool:
        """
        Send notification through all available services, in priority order.
        If one service succeeds, stop and return True.
        """
        for service in self.services:
            try:
                if await service.send_notification(message):
                    return True
            except Exception as e:
                logger.error(f"Error sending notification via service {service.__class__.__name__}: {e}")
        
        return False
    
    def get_cached_notifications(self, limit: int = 100):
        """Get notifications that have been cached"""
        return self.cache_service.get_notifications(limit)

# Initialize notification manager
notification_manager = NotificationManager()

# Function to send notifications to use in API
async def send_notification(message: str) -> bool:
    """Send notification via notification manager"""
    return await notification_manager.send_notification(message)

# Lifespan context manager for managing connections
@asynccontextmanager
async def lifespan(app: FastAPI):
    # Startup logic (database connection handled in db_connector)
    logger.info("Starting MongoDB API server...")
    yield
    # Cleanup on shutdown
    logger.info("Shutting down MongoDB API server...")

app = FastAPI(
    title="Telegram Bot MongoDB API",
    lifespan=lifespan
)

# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Rate limiting counters and cache
request_counts = {}
last_cleanup_time = time.time()
RATE_LIMIT_WINDOW = 60  # 1 minute
RATE_LIMIT_MAX = 50  # 50 requests per minute
cache_data = {}
CACHE_TTL = 60  # 1 minute

# WebSocket connections manager
class ConnectionManager:
    def __init__(self):
        self.active_connections: list[WebSocket] = []

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)
        logger.info(f"New WebSocket connection. Total active: {len(self.active_connections)}")

    def disconnect(self, websocket: WebSocket):
        if websocket in self.active_connections:
            self.active_connections.remove(websocket)
            logger.info(f"WebSocket disconnected. Total active: {len(self.active_connections)}")

    async def broadcast(self, message: dict):
        disconnected = []
        for i, connection in enumerate(self.active_connections):
            try:
                await connection.send_json(message)
            except Exception as e:
                logger.error(f"Error broadcasting message: {e}")
                disconnected.append(i)
                
        # Remove disconnected websockets
        for i in sorted(disconnected, reverse=True):
            if i < len(self.active_connections):
                self.disconnect(self.active_connections[i])

manager = ConnectionManager()

# Session data model
class SessionData(BaseModel):
    session_id: str
    user_id: int
    username: str = None
    first_name: str = None
    last_name: str = None
    timestamp: str = None
    action: str
    message: str = ""
    factor: str

@app.middleware("http")
async def add_process_time_header(request: Request, call_next):
    """Middleware for measuring processing time and rate limiting"""
    # Clean up old rate limit data
    current_time = time.time()
    global last_cleanup_time
    if current_time - last_cleanup_time > RATE_LIMIT_WINDOW:
        request_counts.clear()
        last_cleanup_time = current_time

    # Check rate limit
    client_ip = request.client.host
    current_count = request_counts.get(client_ip, 0)
    if current_count >= RATE_LIMIT_MAX:
        return JSONResponse(
            status_code=429,
            content={"detail": "Rate limit exceeded. Please try again later."}
        )
    
    # Increment count
    request_counts[client_ip] = current_count + 1
    
    # Process request with timing
    start_time = time.time()
    response = await call_next(request)
    process_time = time.time() - start_time
    
    # Add process time header
    response.headers["X-Process-Time"] = str(process_time)
    return response

# Function to process unanswered questions asynchronously
async def process_unanswered_question(session_id: str, background_tasks: BackgroundTasks):
    """
    Process any unanswered question in the background
    """
    async def _process():
        try:
            # Delay a bit to allow the database to update
            await asyncio.sleep(1)
            
            logger.info(f"Processing unanswered question for session {session_id}")
            
            # Get unanswered question by specific session_id
            item = await get_unanswered_question_by_session_id(session_id)
            
            if item:
                logger.info(f"Found complete information for session {session_id}")
                
                # Log without emoji
                log_message = (
                    "Unanswered question detected in background!\n"
                    f"User question: {item['user_question'].get('message', 'No message')}\n"
                    f"RAG response: {item['session'].get('message', 'No message')}"
                )
                logger.info(log_message)
                
                # Ensure data is JSON serializable
                notification_data = {
                    "type": "unanswered_question",
                    "data": item
                }
                
                # Broadcast to all active connections
                if manager.active_connections:
                    await manager.broadcast(notification_data)
                    logger.info(f"Sent notification via WebSocket for session {session_id}")
                else:
                    logger.warning(f"No active WebSocket connections to send notification for session {session_id}")
                    
                    # Try to send direct notification via Telegram if no WebSocket
                    try:
                        # Message for admin with emoji
                        notification_message = (
                            f"🚨 *Unanswered question detected!*\n\n"
                            f"👤 *User:* {item['user_question'].get('first_name', '')} {item['user_question'].get('last_name', '')} (@{item['user_question'].get('username', 'Unknown')})\n"
                            f"💬 *Question:* {item['user_question'].get('message', 'No content')}\n"
                            f"🤖 *RAG response:* {item['session'].get('message', 'No message')}\n"
                            f"🆔 *Session ID:* `{session_id}`\n"
                            f"📝 *Processed by:* Background process"
                        )
                        
                        sent = await send_notification(notification_message)
                        if sent:
                            logger.info(f"Sent direct notification via notification service from background for session {session_id}")
                    except Exception as e:
                        logger.error(f"Error sending notification from background: {e}")
            else:
                logger.warning(f"Could not find complete information for session {session_id}")
                # Try to send direct notification
                try:
                    # Notification with emoji for admin
                    notification_message = (
                        f"🚨 *Unanswered question detected!*\n\n"
                        f"🔍 *RAG responded 'I don't know' for session {session_id}*\n"
                        f"⚠️ *Could not find detailed information about the question*"
                    )
                    
                    sent = await send_notification(notification_message)
                    if sent:
                        logger.info(f"Sent direct notification from background for session {session_id}")
                except Exception as e:
                    logger.error(f"Could not send direct notification from background: {e}")
                
        except Exception as e:
            logger.error(f"Error processing unanswered question: {e}")
            import traceback
            logger.error(traceback.format_exc())
    
    # Add to background tasks
    background_tasks.add_task(_process)

@app.post("/session")
async def create_session(session: SessionData, background_tasks: BackgroundTasks):
    """
    Save a new session chat to the database
    """
    start_time = time.time()
    try:
        # If timestamp is not provided, set it to current time in Vietnam timezone (UTC+7)
        if not session.timestamp:
            session.timestamp = (datetime.utcnow().replace(microsecond=0).isoformat() + "+07:00")
        
        # Ensure all string fields are not None
        if session.username is None:
            session.username = ""
        if session.first_name is None:
            session.first_name = ""
        if session.last_name is None:
            session.last_name = ""
        if session.message is None:
            session.message = ""
            
        # Log details of session being saved
        logger.info(f"Saving session: id={session.session_id}, factor={session.factor}, action={session.action}")
        logger.info(f"Message content: {session.message[:100]}...")
        
        session_data = session.dict()
        session_id = await save_session(session_data)
        
        # Check if this is a RAG response with "I don't know"
        is_rag_unanswered = (
            session.factor == "RAG" and 
            session.message and
            session.message.strip().lower().startswith("i don't know")
        )

        logger.info(f"Is unanswered RAG response: {is_rag_unanswered}")
        
        if is_rag_unanswered:
            # Process unanswered question in background
            background_tasks.add_task(process_unanswered_question, session.session_id, background_tasks)
        
        process_time = time.time() - start_time
        logger.info(f"Session created in {process_time:.2f} seconds: {session_id}")
        return {"status": "success", "session_id": session_id, "process_time": process_time}
    except Exception as e:
        logger.error(f"Error creating session: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/unanswered-questions")
async def fetch_unanswered_questions():
    """
    Fetch sessions where RAG responded with "I don't know"
    """
    cache_key = "unanswered_questions"
    current_time = time.time()
    
    # Check cache first
    if cache_key in cache_data:
        cache_entry = cache_data[cache_key]
        if current_time - cache_entry["timestamp"] < CACHE_TTL:
            logger.info("Returning cached unanswered questions")
            return cache_entry["data"]
    
    try:
        start_time = time.time()
        results = await get_unanswered_questions()
        
        # Cache the results
        cache_data[cache_key] = {
            "data": {"status": "success", "data": results},
            "timestamp": current_time
        }
        
        process_time = time.time() - start_time
        logger.info(f"Fetched unanswered questions in {process_time:.2f} seconds")
        return {"status": "success", "data": results}
    except Exception as e:
        logger.error(f"Error fetching unanswered questions: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    """
    WebSocket endpoint for real-time notifications
    """
    await manager.connect(websocket)
    try:
        while True:
            # Keep connection alive with ping/pong
            try:
                # Wait for messages but don't block too long
                data = await asyncio.wait_for(websocket.receive_text(), timeout=30)
                # Echo back any message received
                await websocket.send_text(f"Echo: {data}")
            except asyncio.TimeoutError:
                # Send a ping to keep the connection alive
                await websocket.send_text("ping")
            except Exception as e:
                logger.error(f"WebSocket error: {e}")
                break
    except WebSocketDisconnect:
        manager.disconnect(websocket)

@app.get("/health")
async def health_check():
    """
    Health check endpoint
    """
    return {
        "status": "healthy", 
        "active_websockets": len(manager.active_connections),
        "cache_entries": len(cache_data)
    }

# Hugging Face Spaces compatibility
# Add a root endpoint for spaces UI
@app.get("/")
async def root():
    """
    Root endpoint for Hugging Face Spaces UI
    """
    return {
        "message": "Telegram Bot MongoDB API is running",
        "docs_url": "/docs",
        "health_check": "/health",
        "endpoints": [
            {"path": "/session", "method": "POST", "description": "Save a new session chat"},
            {"path": "/unanswered-questions", "method": "GET", "description": "Fetch unanswered questions"},
            {"path": "/ws", "method": "WebSocket", "description": "WebSocket for real-time notifications"},
            {"path": "/health", "method": "GET", "description": "Health check endpoint"}
        ]
    }

if __name__ == "__main__":
    import uvicorn
    
    # Get port from environment or default to 7860 (Hugging Face default)
    port = int(os.environ.get("PORT", 7860))
    
    uvicorn.run("app:app", host="0.0.0.0", port=port)