import os

from fastapi import FastAPI
from fastapi import UploadFile
from fastapi import File
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel

from rag import create_vector_store
from rag import ask_question


app = FastAPI(
    title="AI Document QA Bot",
    description="RAG-based Question Answering using LangChain, ChromaDB and Gemini",
    version="1.0.0"
)


UPLOAD_DIRECTORY = "uploads"

os.makedirs(UPLOAD_DIRECTORY, exist_ok=True)


class QuestionRequest(BaseModel):
    question: str


@app.get("/")
def home():
    return {
        "message": "AI Document QA Bot"
    }


@app.post("/upload")
async def upload_pdf(
        file: UploadFile = File(...)
):

    file_path = os.path.join(
        UPLOAD_DIRECTORY,
        file.filename
    )

    with open(file_path, "wb") as f:
        f.write(await file.read())

    create_vector_store(file_path)

    return {
        "message": "PDF uploaded successfully"
    }


@app.post("/ask")
def ask(
        request: QuestionRequest
):

    answer = ask_question(
        request.question
    )

    return {
        "answer": answer
    }