from fastapi import FastAPI, Query
from pydantic import BaseModel
from typing import Any
import pandas as pd
from mojica_agent import MojicaAgent
from config import Config

app = FastAPI()

mojica_bot = MojicaAgent(Config)


# * Esquema de entrada como marshmellow
class QuestionRequest(BaseModel):
    question: str


class AnswerResponse(BaseModel):
    sql: str
    result: Any



@app.post("/")
def ask_question(req: QuestionRequest):
    sql, result = mojica_bot.consult(req.question)

    # Si es dataframe lo convertimos a json
    if isinstance(result, pd.DataFrame):
        result = result.to_json(orient="records")

    return {"sql": sql, "result": result}
    # return {"sql": "WASA"}


# @app.post("/", response_model=AnswerResponse)
# def ask_question(req: QuestionRequest):
#     sql, result = mojica_bot.consult(req.question)

#     # * Si es dataframe lo convertimos a json
#     if isinstance(result, pd.DataFrame):
#         result = result.to_dict(orient="records")

#     return {"sql": sql, "result": result}

# @app.get("/")
# def greet_json():
#     return {"Hello": "World!"}