api/main.py

from fastapi import FastAPI, HTTPException, Depends, Request, status, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from fastapi.security import OAuth2PasswordBearer
from fastapi.middleware.gzip import GZipMiddleware
from typing import Dict, Any, Optional, List
import time
import logging
from datetime import datetime
from pydantic import BaseModel, Field
import os
import asyncio
from tenacity import retry, stop_after_attempt, wait_exponential

from config.config import settings
from core.rag_engine import RAGEngine
from core.user_profile import UserProfile, UserPreferences

# Define missing types
class ChatRequest(BaseModel):
    message: str
    chat_history: Optional[List[Dict[str, str]]] = None

class ChatResponse(BaseModel):
    answer: str
    sources: Optional[List[str]] = None
    suggested_questions: Optional[List[str]] = None

class ErrorResponse(BaseModel):
    error: str
    detail: Optional[str] = None
    timestamp: str = Field(default_factory=lambda: datetime.utcnow().isoformat())
    request_id: Optional[str] = None

class UserProfileResponse(BaseModel):
    profile: Dict[str, Any]

class UserPreferencesUpdate(BaseModel):
    preferences: Dict[str, Any]

# Setup logging with rotation
from logging.handlers import RotatingFileHandler

logging.basicConfig(
    level=getattr(logging, settings.LOG_LEVEL),
    format=settings.LOG_FORMAT,
    handlers=[
        logging.StreamHandler(),
        RotatingFileHandler(
            "api.log",
            maxBytes=10 * 1024 * 1024,  # 10MB
            backupCount=5,
        ),
    ],
)
logger = logging.getLogger(__name__)

app = FastAPI(
    title=settings.PROJECT_NAME,
    description="AI-powered travel assistant API",
    version=settings.VERSION,
    docs_url="/docs",  # Always show docs on HF Spaces
    redoc_url="/redoc",
)

# Add security headers middleware
@app.middleware("http")
async def add_security_headers(request: Request, call_next):
    response = await call_next(request)
    response.headers["X-Content-Type-Options"] = "nosniff"
    response.headers["X-Frame-Options"] = "DENY"
    response.headers["X-XSS-Protection"] = "1; mode=block"
    response.headers["Strict-Transport-Security"] = (
        "max-age=31536000; includeSubDomains"
    )
    return response

# Add CORS middleware with validation
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Allow all origins for Hugging Face Spaces
    allow_credentials=True,
    allow_methods=["GET", "POST", "PUT", "DELETE"],
    allow_headers=["*"],
    max_age=3600,
)

# Add Gzip compression
app.add_middleware(GZipMiddleware, minimum_size=1000)

# Initialize core components with retry
@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10))
async def initialize_components():
    try:
        global rag_engine, user_profile
        rag_engine = RAGEngine()
        user_profile = UserProfile()
        logger.info("Core components initialized successfully")
    except Exception as e:
        logger.error(f"Failed to initialize core components: {str(e)}", exc_info=True)
        raise

# Initialize components asynchronously
asyncio.create_task(initialize_components())

# OAuth2 scheme for token authentication
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")

from api.dependencies import (
    get_current_user,
    rate_limit,
    cleanup,
)

@app.exception_handler(Exception)
async def global_exception_handler(request: Request, exc: Exception):
    """Global exception handler with request ID"""
    request_id = request.headers.get("X-Request-ID", "unknown")
    logger.error(
        f"Unhandled exception: {str(exc)}",
        exc_info=True,
        extra={"request_id": request_id},
    )
    return JSONResponse(
        status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
        content=ErrorResponse(
            error="Internal Server Error", detail=str(exc), request_id=request_id
        ).dict(),
    )

@app.middleware("http")
async def add_process_time_header(request: Request, call_next):
    """Add processing time header to response"""
    start_time = time.time()
    try:
        response = await call_next(request)
        process_time = time.time() - start_time
        response.headers["X-Process-Time"] = str(process_time)
        return response
    except Exception as e:
        logger.error(f"Error in middleware: {str(e)}", exc_info=True)
        raise

@app.get("/")
async def root():
    """Root endpoint with version info"""
    return {
        "message": "Welcome to TravelMate AI Assistant API",
        "version": settings.VERSION,
        "environment": settings.DEBUG,  # Use DEBUG setting for environment
    }

@app.post(
    "/chat",
    response_model=ChatResponse,
    responses={
        400: {"model": ErrorResponse},
        401: {"model": ErrorResponse},
        429: {"model": ErrorResponse},
        500: {"model": ErrorResponse},
    },
)
@rate_limit
async def chat(
    request: ChatRequest,
    background_tasks: BackgroundTasks,
    current_user: Dict[str, Any] = Depends(get_current_user),
):
    """Process chat request with enhanced validation"""
    try:
        # Validate request size
        if len(request.message) > settings.MAX_MESSAGE_LENGTH:
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=f"Message too long. Maximum length is {settings.MAX_MESSAGE_LENGTH} characters",
            )

        # Validate chat history
        if request.chat_history:
            if len(request.chat_history) > settings.MAX_CHAT_HISTORY:
                raise HTTPException(
                    status_code=status.HTTP_400_BAD_REQUEST,
                    detail=f"Chat history too long. Maximum length is {settings.MAX_CHAT_HISTORY} messages",
                )
            for msg in request.chat_history:
                if not isinstance(msg, dict) or not all(
                    k in msg for k in ["user", "assistant"]
                ):
                    raise HTTPException(
                        status_code=status.HTTP_400_BAD_REQUEST,
                        detail="Invalid chat history format",
                    )

        # Process query with RAG engine
        result = await asyncio.wait_for(
            rag_engine.process_query(
                query=request.message,
                chat_history=request.chat_history,
                user_id=current_user["user_id"],
            ),
            timeout=settings.QUERY_TIMEOUT,
        )

        # Add cleanup task
        background_tasks.add_task(cleanup)

        return ChatResponse(
            answer=result["answer"],
            sources=result.get("metadata", {}).get("sources", []),
            suggested_questions=result.get("suggested_questions", []),
        )
    except asyncio.TimeoutError:
        raise HTTPException(
            status_code=status.HTTP_504_GATEWAY_TIMEOUT, detail="Request timed out"
        )
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Error processing chat request: {str(e)}", exc_info=True)
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail="Error processing chat request",
        )

@app.get(
    "/profile",
    response_model=UserProfileResponse,
    responses={401: {"model": ErrorResponse}, 500: {"model": ErrorResponse}},
)
async def get_profile(current_user: Dict[str, Any] = Depends(get_current_user)):
    """Get user profile with enhanced error handling"""
    try:
        profile = await asyncio.wait_for(
            user_profile.get_profile(current_user["user_id"]),
            timeout=settings.PROFILE_TIMEOUT,
        )
        if not profile:
            raise HTTPException(
                status_code=status.HTTP_404_NOT_FOUND, detail="Profile not found"
            )
        return UserProfileResponse(**profile)
    except asyncio.TimeoutError:
        raise HTTPException(
            status_code=status.HTTP_504_GATEWAY_TIMEOUT, detail="Request timed out"
        )
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Error getting user profile: {str(e)}", exc_info=True)
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail="Error retrieving profile",
        )

@app.put(
    "/profile/preferences",
    responses={
        400: {"model": ErrorResponse},
        401: {"model": ErrorResponse},
        500: {"model": ErrorResponse},
    },
)
async def update_preferences(
    preferences: UserPreferencesUpdate,
    current_user: Dict[str, Any] = Depends(get_current_user),
):
    """Update user preferences with validation"""
    try:
        # Validate preferences
        try:
            UserPreferences(**preferences.preferences)
        except Exception as e:
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail=f"Invalid preferences: {str(e)}",
            )

        success = await asyncio.wait_for(
            user_profile.update_profile(
                current_user["user_id"], {"preferences": preferences.preferences}
            ),
            timeout=settings.PROFILE_TIMEOUT,
        )

        if not success:
            raise HTTPException(
                status_code=status.HTTP_400_BAD_REQUEST,
                detail="Failed to update preferences",
            )
        return {"message": "Preferences updated successfully"}
    except asyncio.TimeoutError:
        raise HTTPException(
            status_code=status.HTTP_504_GATEWAY_TIMEOUT, detail="Request timed out"
        )
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Error updating preferences: {str(e)}", exc_info=True)
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail="Error updating preferences",
        )

@app.get("/health", responses={500: {"model": ErrorResponse}})
async def health_check():
    """Health check endpoint with detailed status"""
    try:
        # Check core components
        if not rag_engine or not user_profile:
            raise HTTPException(
                status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
                detail="Core components not initialized",
            )

        return {
            "status": "healthy",
            "timestamp": datetime.utcnow().isoformat(),
            "version": settings.VERSION,
            "environment": settings.DEBUG,  # Use DEBUG setting for environment
            "components": {
                "rag_engine": "ok",
                "user_profile": "ok",
            },
        }
    except Exception as e:
        logger.error(f"Health check failed: {str(e)}", exc_info=True)
        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="Service unhealthy"
        )

@app.on_event("shutdown")
async def shutdown_event():
    """Cleanup on shutdown"""
    await cleanup()

if __name__ == "__main__":
    import uvicorn

    uvicorn.run(
        "api.main:app",
        host="0.0.0.0",
        port=int(os.getenv("PORT", 7860)),
        reload=False,  # Set reload to False for production
    )