import PIL
from fastapi import FastAPI, File, UploadFile
from pydantic import BaseModel
from fastapi.responses import JSONResponse
from utils.model_func import class_id_to_label, load_model, transform_image

model = None
app = FastAPI()


class ImageClass(BaseModel):
    prediction: str

class TextClass(BaseModel):
    text: str


@app.on_event("startup")
async def startup_event():
    global model
    # Здесь используйте функцию из utils.model_func для загрузки модели
    model = load_model()


# @app.post('/classify')
# async def classify_image(file: UploadFile = File(...)):
#     # Здесь используйте функцию из utils.model_func для классификации изображения
#     image_bytes = await file.read()
#     prediction = transform_image(image_bytes, model)
#     return {"prediction": prediction}

@app.post('/classify')
async def classify_image(file: UploadFile = File(...)):
    # Use the function from utils.model_func to classify the image
    image = PIL.Image.open(file.file)
    adapted_image = transform_image(image)
    pred_index = model(adapted_image.unsqeeze(0).detach().cpu().numpy().argmax())
    imagenet_class = class_id_to_label(pred_index)     
    response = ImageClass(prediction=imagenet_class)   
    return response

          
                       

@app.post('/clf_text')
async def classify_text(text_data: TextClass):
    # Здесь используйте функцию из utils.model_func для классификации текста
    text = text_data.text
    prediction = class_id_to_label(text, model)
    return {"prediction": prediction}