app/main.py

import logging
from contextlib import asynccontextmanager
from pathlib import Path

from dotenv import load_dotenv
from fastapi import FastAPI

from app.api.routes import router
from app.core.config import get_settings
from app.dependencies import set_rag_service
from app.services.rag_service import RAGService

logger = logging.getLogger(__name__)
load_dotenv()


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Manage application lifecycle - load index on startup"""
    settings = get_settings()

    logger.info("Initializing RAG service...")
    rag_service = RAGService(settings)
    set_rag_service(rag_service)
    index_path = Path(settings.index_path)
    colbert_index_path = Path(settings.colbert_index_path)

    index_exists = (
        (index_path / "bm25_index.pkl").exists() if not settings.colbert_only else True
    )
    colbert_exists = (
        (colbert_index_path / "documents.pkl").exists()
        if settings.use_colbert
        else True
    )

    if index_exists and colbert_exists:
        try:
            rag_service.load_index()
            logger.info("Existing RAG index loaded successfully")
        except Exception as e:
            logger.error(f"Failed to load index: {e}")
            logger.info("Will attempt to rebuild index...")
            index_exists = False
    else:
        logger.info("No existing index found")

    if not (index_exists and colbert_exists):
        markdown_path = Path(settings.markdown_directory)
        if markdown_path.exists() and any(markdown_path.glob("**/*.md")):
            try:
                logger.info(f"Building new index from {settings.markdown_directory}...")
                rag_service.build_index(settings.markdown_directory)
                logger.info("Index built successfully")
            except Exception as e:
                logger.error(f"Failed to build index: {e}")
                logger.warning(
                    "Application starting without index - you'll need to build it manually"
                )
        else:
            logger.warning(f"No markdown files found in {settings.markdown_directory}")
            logger.warning(
                "Application starting without index - add markdown files and restart"
            )

    yield

    logger.info("Shutting down...")


app = FastAPI(
    title="RAG API",
    description="Retrieval Augmented Generation API for Document Q&A",
    version="1.0.0",
    lifespan=lifespan,
)

app.include_router(router, prefix="/api/v1")