# app.py
from fastapi import FastAPI, UploadFile, File, Body
from fastapi.responses import JSONResponse
import numpy as np
from PIL import Image
import io
import tensorflow as tf

app = FastAPI()

class Service:
    def __init__(self):
        self.model = None
        self.models={}
        self.active_model = "default"

    def activate_model(self, name):

        self.model = self.models[name]
        self.active_model = name
        return 0

    def load_model(self, path, name):
        self.models[name] = tf.keras.models.load_model(path)
        return 0

    def get_accept_size(self):
        print(self.model.layers[0].input)

        return list(self.model.input_shape)

    def get_output_size(self):
        return list(self.model.layers[-1].output[0].shape[1:4])


    def predict(self, data: bytes):
        target_size = self.get_accept_size()
        image = np.reshape(
            np.frombuffer(data, dtype=np.uint8),
            (target_size[1], target_size[2], target_size[3])
        ) / 255.0

        predictions = self.model.predict(np.array([image]))

        if predictions[0].ndim == 2:
            classed = predictions[0]
        else:
            classed = tf.argmax(predictions[0], axis=2)

        return JSONResponse({"result": classed.numpy().reshape(-1).tolist()})

    def predict_raw(self, data):
        target_size = self.get_accept_size()
        image = np.reshape(np.frombuffer(data, dtype=np.uint8), (target_size[1], target_size[2], target_size[3])) / 255.

        predictions = self.model(np.array([image])).numpy()

        if predictions[0].ndim==2:
            classed = predictions[0][:,:,np.newaxis].clip(0,255)
            return (classed).astype(np.uint8).reshape(-1).tolist()
        else:
            classed = predictions[0][:,:,:]

        return (classed*255).astype(np.uint8).reshape(-1).tolist()

svc = Service()

def _prep_raw_bytes_for_predict(buf: bytes) -> bytes:
    """Decode JPEG/PNG -> RGB, resize to model size, return raw uint8 bytes HxWxC."""
    h, w, c = svc.get_accept_size()[1:4]
    img = Image.open(io.BytesIO(buf)).convert("RGB")
    if img.size != (w, h):
        img = img.resize((w, h), Image.BILINEAR)
    arr = np.asarray(img, dtype=np.uint8)  # HxWx3 uint8
    return arr.tobytes()

@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    buf = await file.read()
    raw = _prep_raw_bytes_for_predict(buf)
    return svc.predict(raw)

@app.post("/predict_raw")
async def predict_raw(file: UploadFile = File(...)):
    buf = await file.read()
    raw = _prep_raw_bytes_for_predict(buf)
    return svc.predict_raw(raw)

@app.post("/load_model")
async def load_model(path: str = Body(...), name: str = Body(...)):
    try:
        svc.load_model(path, name)
        return JSONResponse({"status": "success", "message": f"Model {name} loaded successfully."})
    except Exception as e:
        return JSONResponse({"status": "error", "message": str(e)}, status_code=500)

@app.post("/activate_model")
async def activate_model(name: str = Body(...)):
    try:
        svc.activate_model(name)
        return JSONResponse({"status": "success", "message": f"Model {name} activated successfully."})
    except Exception as e:
        return JSONResponse({"status": "error", "message": str(e)}, status_code=500)