from fastapi import FastAPI, HTTPException
from omegaconf import OmegaConf
from pydantic import BaseModel

from rag_qa import RagQA

CONFIG_PATH = "rag_pipeline/conf/inference.yaml"

app = FastAPI()
rag_qa: RagQA = None
is_rag_qa_ready: bool = False


class DiscoveredSource(BaseModel):
    name: str
    text: str
    index_id: int


class Request(BaseModel):
    question: str


class Response(BaseModel):
    answer: str
    sources: list[DiscoveredSource]


@app.on_event("startup")
async def startup_event() -> None:
    global rag_qa, is_rag_qa_ready
    conf = OmegaConf.load(CONFIG_PATH)
    rag_qa = RagQA(conf)
    print("Server startup: Initializing RagQA")
    await rag_qa.load()
    is_rag_qa_ready = True
    print("Server startup: Initialized RagQA.")


@app.post("/answer", response_model=Response)
async def answer(request: Request) -> Response:
    if not is_rag_qa_ready:
        raise HTTPException(
            status_code=503,
            detail="Service unavailable: RagQA is still initializing. Please try again later.",
        )

    answer, sources = rag_qa.answer(request.question)
    discovered_sources = [DiscoveredSource(**source.__dict__) for source in sources]
    return Response(answer=answer, sources=discovered_sources)