# app.py
from fastapi import FastAPI, Response, UploadFile, File, Body, HTTPException
from pydantic import BaseModel
from typing import List
from fastapi.responses import JSONResponse
import numpy as np
from PIL import Image
import io
import tensorflow as tf
import time
app = FastAPI()

# ==== Schemas (show up in Swagger) ====
class PredictOut(BaseModel):
    result: List[int]
    model_config = {"json_schema_extra": {"examples": [{"result": [0,1,2,3]}]}}

class StatusOut(BaseModel):
    status: str
    message: str

class AcceptSizeOut(BaseModel):
    status: str
    accept_size: List[int]

class LoadModelIn(BaseModel):
    path: str
    name: str

class ActivateModelIn(BaseModel):
    name: str

# ==== Service ====
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):
        if name in self.models:
            # Model already loaded, skip loading
            return 0
        self.models[name] = tf.keras.models.load_model(path)
        return 0

    def get_accept_size(self):
        print(self.model.input_shape)
        return list(self.model.input_shape)  # (B,H,W,C)

    def get_output_size(self):
        return list(self.model.output_shape)  # (B,H,W,C) or (B,H,W)

    def predict(self, data: bytes) -> List[int]:
        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 classed.numpy().reshape(-1).tolist()

    def predict_raw(self, data: bytes) -> np.ndarray:
        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({"input_layer": np.array([image], dtype=np.float32)}).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][:, :, :]

        res = (classed * 255).astype(np.uint8).reshape(-1)
        return np.asarray(res)

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()

# ==== Routes ====
@app.post("/predict", response_model=bytes, summary="Predict (classified indices)")
async def predict(file: UploadFile = File(...)):
    try:
        buf = await file.read()
        raw = _prep_raw_bytes_for_predict(buf)
        result = svc.predict(raw)
        return Response(content=result, media_type="application/octet-stream")
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))

@app.post("/predictRaw", response_model=bytes, summary="Predict (raw/uint8 flattened)")
async def predict_raw(file: UploadFile = File(...)):
    try:
        start_time = time.time()
        buf = await file.read()
        result = svc.predict_raw(buf)
        elapsed_time = time.time() - start_time
        print(f"Prediction took {elapsed_time:.2f} seconds")
        return Response(content=result.tobytes(), media_type="application/octet-stream")
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))

@app.post("/loadModel", response_model=StatusOut, summary="Load a model from path")
async def load_model(payload: LoadModelIn):
    try:
        svc.load_model(payload.path, payload.name)
        return StatusOut(status="success", message=f"Model {payload.name} loaded successfully.")
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/activateModel", response_model=StatusOut, summary="Activate a loaded model")
async def activate_model(payload: ActivateModelIn):
    try:
        svc.activate_model(payload.name)
        return StatusOut(status="success", message=f"Model {payload.name} activated successfully.")
    except KeyError:
        raise HTTPException(status_code=404, detail=f"Model {payload.name} not found.")
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/getAcceptSize", response_model=AcceptSizeOut, summary="Get input tensor shape")
async def get_accept_size():
    try:
        size = svc.get_accept_size()[1:]  # Exclude batch size
        return AcceptSizeOut(status="success", accept_size=size)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/getOutputSize", response_model=AcceptSizeOut, summary="Get output tensor shape")
async def get_output_size():
    try:
        size = svc.get_output_size()[1:]
        return AcceptSizeOut(status="success", accept_size=size)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))