from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from fastapi.staticfiles import StaticFiles
from services.generate_questions_service import GenerateQuestionsService
from typing import Optional

# from data.load_data import retriever_pre

generate_questions_service = GenerateQuestionsService()


class Body(BaseModel):
    school_subject: Optional[str] = None
    subject: Optional[str] = None
    difficultie: Optional[str] = None
    promptValue: Optional[str] = None


app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.post("/generate_questions")
def generate_questions(body: Body):
    if body.promptValue:
        query = body.promptValue
    else:
        school_subject = body.school_subject
        subject = body.subject
        difficultie = body.difficultie
        query = f"Quero que você gere questões de {school_subject}, sendo do assunto: {subject} e sendo da dificuldade: {difficultie}."
    res = generate_questions_service.handle(f"""{query}""")
    return res


app.mount("/", StaticFiles(directory="static", html=True), name="static")

if __name__ == "__main__":
    import uvicorn

    uvicorn.run("app:app", host="0.0.0.0", port=8000)