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


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

        # For server
        self.fhemodel_server = FHEModelServer(path + "/compiled_model")

    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 = data.pop("encrypted_inputs", data)

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

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

        return encrypted_prediction