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


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


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


@app.get("/")
def hello():
    return {"massage": "we are running!"}


@app.get("/generate/")
def generate(prompt: str, lang: str):
    # request_data = prompt.dict()

    # text = request_data["prompts"]

    if lang:
        generated_recepie = generation_function(
            prompt, lang=lang)
    else:
        generated_recepie = generation_function(prompt)

    return generated_recepie