from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from generator import generation_function
from pydantic import BaseModel


app = FastAPI(docs_url="/")
app.add_middleware(
    CORSMiddleware,
    allow_origins=['*'],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


class Prompt(BaseModel):
    prompts: str
    lang: str


@app.get("/test")
def test(input: str):
    output = input
    return {"input": output}


@app.post("/generate/")
def generate(prompt: Prompt):
    request_data = prompt.dict()

    text = request_data["prompts"]

    if request_data.get("lang"):
        generated_recepie = generation_function(
            text, lang=request_data["lang"])
    else:
        generated_recepie = generation_function(text)

    return generated_recepie