from typing import Dict, List, Any
import numpy as np
from concrete.ml.deployment import FHEModelServer


def from_json(python_object):
    if "__class__" in python_object:
        return bytes(python_object["__value__"])


def to_json(python_object):
    if isinstance(python_object, bytes):
        return {"__class__": "bytes", "__value__": list(python_object)}
    # if isinstance(python_object, bytes):
    #     return list(python_object)
    raise TypeError(repr(python_object) + " is not JSON serializable")


class EndpointHandler:
    def __init__(self, path="benoit.model"):

        # For server
        self.fhemodel_server = FHEModelServer(path)

    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        """
         data args:
              inputs (:obj: `str`)
              date (:obj: `str`)
        Return:
              A :obj:`list` | `dict`: will be serialized and returned
        """

        # Get inputs
        encrypted_inputs = from_json(data.pop("encrypted_inputs", data))

        # Get keys
        evaluation_keys = from_json(data.pop("evaluation_keys", data))

        # Run CML prediction
        encrypted_prediction = self.fhemodel_server.run(encrypted_inputs, evaluation_keys)

        return to_json(encrypted_prediction)