import datetime
import pprint
from typing import Any, Dict, Tuple, Optional, List, Union

import gql  # type: ignore
import numpy
from .inference_decode import decode_inference_result, to_nd_array_list
from .object import InferenceError


class InferenceResult(object):
    def __init__(self, gql_client: Optional[gql.Client], data: Dict[str, Any]) -> None:
        """Initializes an InferenceResult.

        :param gql.Client gql_client: GQL client that this can pass to created objects.
        :param Dict[str, Any] data: Response parsed from JSON inference result body.
        """
        if "error" in data:
            raise InferenceError(data)

        self._model_version = (data["model_name"], data["model_version"])
        self._input_data = data["original_data"]
        # Assumes that timestamp comes back in integer milliseconds since epoch
        self._timestamp = datetime.datetime.fromtimestamp(data["time"] / 1000)
        # Assumes that elapsed time comes back in integer nanoseconds
        elapsed = data["elapsed"]
        elapsed = sum(elapsed) if isinstance(elapsed, list) else elapsed
        self._time_elapsed = datetime.timedelta(microseconds=elapsed / 1000)

        # TODO: we may want to match on key types and special case them, but
        # this handles all types in a reasonable fashion (i.e. exactly the
        # same)
        self._data: List[numpy.ndarray] = to_nd_array_list(
            decode_inference_result(data)
        )
        if "shadow_data" in data:
            self._shadow_data: Union[Dict[str, numpy.ndarray], None] = {
                shadow_model: sd for shadow_model, sd in data["shadow_data"].items()
            }
        else:
            self._shadow_data = None
        self.raw = data

    def data(self) -> List[numpy.ndarray]:
        """Returns the inference result data."""
        return self._data

    def model_version(self) -> Tuple[str, str]:
        """Returns the model this inference was generated by."""
        # Model is currently passed back as model class + model name, which will
        # necessitate passing the client through the deployment to here in order
        # to call `model_by_name`. Better is to pass the DB surrogate key, which
        # we can construct the object off of.
        #
        # For now, just return the two names, and the caller can use a client to
        # turn this back into a Model if they wish.
        return self._model_version

    def time_elapsed(self) -> datetime.timedelta:
        """Returns the total time taken by the engine."""
        return self._time_elapsed

    def timestamp(self) -> datetime.datetime:
        """Returns the time at which this inference occurred."""
        return self._timestamp

    def input_data(self) -> Dict[str, Any]:
        """Returns the input data for this inference result."""
        return self._input_data

    def __repr__(self) -> str:
        return f"InferenceResult({pprint.pformat(self.raw)})"

    def __str__(self) -> str:
        return repr(self)

    def shadow_data(self) -> Union[Dict[str, numpy.ndarray], None]:
        return self._shadow_data
