import io
from typing import Any, Union

from PIL import Image

from clearml import StorageManager


# Notice Preprocess class Must be named "Preprocess"
class Preprocess(object):
    def __init__(self):
        # set internal state, this will be called only once. (i.e. not per request)
        pass

    def preprocess(self, body: Union[bytes, dict], state: dict, collect_custom_statistics_fn=None) -> Any:
        # we expect to get two valid on the dict x0, and x1
        if isinstance(body, bytes):
            # we expect to get a stream of encoded image bytes
            try:
                Image.open(io.BytesIO(body)).convert("RGB")
            except Exception:
                # value error would return 404, we want to return 500 so any other exception
                raise RuntimeError("Image could not be decoded")
            raise ValueError("send JSON with 'url' (or 'image_url') and 'xywh'")

        if isinstance(body, dict) and ("url" in body.keys() or "image_url" in body.keys()):
            # image is given as url, and is fetched
            url = body.get("image_url") or body.get("url")
            local_file = StorageManager.get_local_copy(remote_url=url)
            xywh = body.get("xywh")
            if not xywh:
                raise ValueError("body must include 'xywh' as 'x,y,w,h'")
            # PersonAttrService.predict 에 전달 (로컬 캐시 경로 + bbox)
            return {"image_url": local_file, "xywh": str(xywh).strip()}

        raise ValueError("body must include 'url' (or 'image_url') and 'xywh'")

    def postprocess(self, data: Any, state: dict, collect_custom_statistics_fn=None) -> dict:
        # post process the data returned from the model inference engine
        # data is the return value from model.predict we will put is inside a return value as Y
        if not isinstance(data, dict):
            # this should not happen
            return dict(output={}, status="FAIL")

        # data is returned as attribute dict from PersonAttrService.predict
        if "output" in data and "status" in data:
            return data
        return dict(output=data, status="PASS")