import inspect
import itertools
import logging
import os
import traceback
from abc import ABC
from typing import Any, Dict, Iterator, List

from clarifai_grpc.grpc.api import resources_pb2, service_pb2
from clarifai_grpc.grpc.api.status import status_code_pb2, status_pb2
from google.protobuf import json_format

from clarifai.runners.utils import data_types
from clarifai.runners.utils.data_utils import DataConverter
from clarifai.runners.utils.method_signatures import (build_function_signature, deserialize,
                                                      get_stream_from_signature, serialize,
                                                      signatures_to_json)

_METHOD_INFO_ATTR = '_cf_method_info'

_RAISE_EXCEPTIONS = os.getenv("RAISE_EXCEPTIONS", "false").lower() in ("true", "1")


class ModelClass(ABC):
  '''
  Base class for model classes that can be run as a service.

  Define predict, generate, or stream methods using the @ModelClass.method decorator.

  Example:

    from clarifai.runners.model_class import ModelClass
    from clarifai.runners.utils.data_types import NamedFields
    from typing import List, Iterator

    class MyModel(ModelClass):

      @ModelClass.method
      def predict(self, x: str, y: int) -> List[str]:
        return [x] * y

      @ModelClass.method
      def generate(self, x: str, y: int) -> Iterator[str]:
        for i in range(y):
          yield x + str(i)

      @ModelClass.method
      def stream(self, input_stream: Iterator[NamedFields(x=str, y=int)]) -> Iterator[str]:
        for item in input_stream:
          yield item.x + ' ' + str(item.y)
  '''

  @staticmethod
  def method(func):
    setattr(func, _METHOD_INFO_ATTR, _MethodInfo(func))
    return func

  def set_output_context(self, prompt_tokens=None, completion_tokens=None):
    """This is used to set the prompt and completion tokens in the Output proto"""
    self._prompt_tokens = prompt_tokens
    self._completion_tokens = completion_tokens

  def load_model(self):
    """Load the model."""

  def _handle_get_signatures_request(self) -> service_pb2.MultiOutputResponse:
    methods = self._get_method_info()
    signatures = {method.name: method.signature for method in methods.values()}
    resp = service_pb2.MultiOutputResponse(status=status_pb2.Status(code=status_code_pb2.SUCCESS))
    output = resp.outputs.add()
    output.status.code = status_code_pb2.SUCCESS
    output.data.text.raw = signatures_to_json(signatures)
    return resp

  def _batch_predict(self, method, inputs: List[Dict[str, Any]]) -> List[Any]:
    """Batch predict method for multiple inputs."""
    outputs = []
    for input in inputs:
      output = method(**input)
      outputs.append(output)
    return outputs

  def _batch_generate(self, method, inputs: List[Dict[str, Any]]) -> Iterator[List[Any]]:
    """Batch generate method for multiple inputs."""
    generators = [method(**input) for input in inputs]
    for outputs in itertools.zip_longest(*generators):
      yield outputs

  def predict_wrapper(
      self, request: service_pb2.PostModelOutputsRequest) -> service_pb2.MultiOutputResponse:
    outputs = []
    try:
      # TODO add method name field to proto
      method_name = 'predict'
      inference_params = get_inference_params(request)
      if len(request.inputs) > 0 and '_method_name' in request.inputs[0].data.metadata:
        method_name = request.inputs[0].data.metadata['_method_name']
      if method_name == '_GET_SIGNATURES':  # special case to fetch signatures, TODO add endpoint for this
        return self._handle_get_signatures_request()
      if method_name not in self._get_method_info():
        raise ValueError(f"Method {method_name} not found in model class")
      method = getattr(self, method_name)
      method_info = method._cf_method_info
      signature = method_info.signature
      python_param_types = method_info.python_param_types
      for input in request.inputs:
        # check if input is in old format
        is_convert = DataConverter.is_old_format(input.data)
        if is_convert:
          # convert to new format
          new_data = DataConverter.convert_input_data_to_new_format(input.data,
                                                                    signature.input_fields)
          input.data.CopyFrom(new_data)
      # convert inputs to python types
      inputs = self._convert_input_protos_to_python(request.inputs, inference_params,
                                                    signature.input_fields, python_param_types)
      if len(inputs) == 1:
        inputs = inputs[0]
        output = method(**inputs)
        outputs.append(
            self._convert_output_to_proto(
                output, signature.output_fields, convert_old_format=is_convert))
      else:
        outputs = self._batch_predict(method, inputs)
        outputs = [
            self._convert_output_to_proto(
                output, signature.output_fields, convert_old_format=is_convert)
            for output in outputs
        ]

      return service_pb2.MultiOutputResponse(
          outputs=outputs, status=status_pb2.Status(code=status_code_pb2.SUCCESS))
    except Exception as e:
      if _RAISE_EXCEPTIONS:
        raise
      logging.exception("Error in predict")
      return service_pb2.MultiOutputResponse(status=status_pb2.Status(
          code=status_code_pb2.FAILURE,
          details=str(e),
          stack_trace=traceback.format_exc().split('\n')))

  def generate_wrapper(self, request: service_pb2.PostModelOutputsRequest
                      ) -> Iterator[service_pb2.MultiOutputResponse]:
    try:
      method_name = 'generate'
      inference_params = get_inference_params(request)
      if len(request.inputs) > 0 and '_method_name' in request.inputs[0].data.metadata:
        method_name = request.inputs[0].data.metadata['_method_name']
      method = getattr(self, method_name)
      method_info = method._cf_method_info
      signature = method_info.signature
      python_param_types = method_info.python_param_types
      for input in request.inputs:
        # check if input is in old format
        is_convert = DataConverter.is_old_format(input.data)
        if is_convert:
          # convert to new format
          new_data = DataConverter.convert_input_data_to_new_format(input.data,
                                                                    signature.input_fields)
          input.data.CopyFrom(new_data)
      inputs = self._convert_input_protos_to_python(request.inputs, inference_params,
                                                    signature.input_fields, python_param_types)
      if len(inputs) == 1:
        inputs = inputs[0]
        for output in method(**inputs):
          resp = service_pb2.MultiOutputResponse()
          self._convert_output_to_proto(
              output,
              signature.output_fields,
              proto=resp.outputs.add(),
              convert_old_format=is_convert)
          resp.status.code = status_code_pb2.SUCCESS
          yield resp
      else:
        for outputs in self._batch_generate(method, inputs):
          resp = service_pb2.MultiOutputResponse()
          for output in outputs:
            self._convert_output_to_proto(
                output,
                signature.output_fields,
                proto=resp.outputs.add(),
                convert_old_format=is_convert)
          resp.status.code = status_code_pb2.SUCCESS
          yield resp
    except Exception as e:
      if _RAISE_EXCEPTIONS:
        raise
      logging.exception("Error in generate")
      yield service_pb2.MultiOutputResponse(status=status_pb2.Status(
          code=status_code_pb2.FAILURE,
          details=str(e),
          stack_trace=traceback.format_exc().split('\n')))

  def stream_wrapper(self, request_iterator: Iterator[service_pb2.PostModelOutputsRequest]
                    ) -> Iterator[service_pb2.MultiOutputResponse]:
    try:
      request = next(request_iterator)  # get first request to determine method
      assert len(request.inputs) == 1, "Streaming requires exactly one input"

      method_name = 'stream'
      inference_params = get_inference_params(request)
      if len(request.inputs) > 0 and '_method_name' in request.inputs[0].data.metadata:
        method_name = request.inputs[0].data.metadata['_method_name']
      method = getattr(self, method_name)
      method_info = method._cf_method_info
      signature = method_info.signature
      python_param_types = method_info.python_param_types

      # find the streaming vars in the signature
      stream_sig = get_stream_from_signature(signature.input_fields)
      if stream_sig is None:
        raise ValueError("Streaming method must have a Stream input")
      stream_argname = stream_sig.name

      for input in request.inputs:
        # check if input is in old format
        is_convert = DataConverter.is_old_format(input.data)
        if is_convert:
          # convert to new format
          new_data = DataConverter.convert_input_data_to_new_format(input.data,
                                                                    signature.input_fields)
          input.data.CopyFrom(new_data)
      # convert all inputs for the first request, including the first stream value
      inputs = self._convert_input_protos_to_python(request.inputs, inference_params,
                                                    signature.input_fields, python_param_types)
      kwargs = inputs[0]

      # first streaming item
      first_item = kwargs.pop(stream_argname)

      # streaming generator
      def InputStream():
        yield first_item
        # subsequent streaming items contain only the streaming input
        for request in request_iterator:
          item = self._convert_input_protos_to_python(request.inputs, inference_params,
                                                      [stream_sig], python_param_types)
          item = item[0][stream_argname]
          yield item

      # add stream generator back to the input kwargs
      kwargs[stream_argname] = InputStream()

      for output in method(**kwargs):
        resp = service_pb2.MultiOutputResponse()
        self._convert_output_to_proto(
            output,
            signature.output_fields,
            proto=resp.outputs.add(),
            convert_old_format=is_convert)
        resp.status.code = status_code_pb2.SUCCESS
        yield resp
    except Exception as e:
      if _RAISE_EXCEPTIONS:
        raise
      logging.exception("Error in stream")
      yield service_pb2.MultiOutputResponse(status=status_pb2.Status(
          code=status_code_pb2.FAILURE,
          details=str(e),
          stack_trace=traceback.format_exc().split('\n')))

  def _convert_input_protos_to_python(self, inputs: List[resources_pb2.Input],
                                      inference_params: dict,
                                      variables_signature: List[resources_pb2.ModelTypeField],
                                      python_param_types) -> List[Dict[str, Any]]:
    result = []
    for input in inputs:
      kwargs = deserialize(input.data, variables_signature, inference_params)
      # dynamic cast to annotated types
      for k, v in kwargs.items():
        if k not in python_param_types:
          continue

        if hasattr(python_param_types[k], "__args__") and getattr(python_param_types[k],
                                                                  "__origin__", None) == Iterator:
          # get the type of the items in the stream
          stream_type = python_param_types[k].__args__[0]

          kwargs[k] = data_types.cast(v, stream_type)
        else:
          kwargs[k] = data_types.cast(v, python_param_types[k])
      result.append(kwargs)
    return result

  def _convert_output_to_proto(self,
                               output: Any,
                               variables_signature: List[resources_pb2.ModelTypeField],
                               proto=None,
                               convert_old_format=False) -> resources_pb2.Output:
    if proto is None:
      proto = resources_pb2.Output()
    serialize({'return': output}, variables_signature, proto.data, is_output=True)
    if convert_old_format:
      # convert to old format
      data = DataConverter.convert_output_data_to_old_format(proto.data)
      proto.data.CopyFrom(data)
    proto.status.code = status_code_pb2.SUCCESS
    if hasattr(self, "_prompt_tokens") and self._prompt_tokens is not None:
      proto.prompt_tokens = self._prompt_tokens
    if hasattr(self, "_completion_tokens") and self._completion_tokens is not None:
      proto.completion_tokens = self._completion_tokens
    self._prompt_tokens = None
    self._completion_tokens = None
    return proto

  @classmethod
  def _register_model_methods(cls):
    # go up the class hierarchy to find all decorated methods, and add to registry of current class
    methods = {}
    for base in reversed(cls.__mro__):
      for name, method in base.__dict__.items():
        method_info = getattr(method, _METHOD_INFO_ATTR, None)
        if not method_info:  # regular function, not a model method
          continue
        methods[name] = method_info
    # check for generic predict(request) -> response, etc. methods
    #for name in ('predict', 'generate', 'stream'):
    #  if hasattr(cls, name):
    #    method = getattr(cls, name)
    #    if not hasattr(method, _METHOD_INFO_ATTR):  # not already put in registry
    #      methods[name] = _MethodInfo(method)
    # set method table for this class in the registry
    return methods

  @classmethod
  def _get_method_info(cls, func_name=None):
    if not hasattr(cls, _METHOD_INFO_ATTR):
      setattr(cls, _METHOD_INFO_ATTR, cls._register_model_methods())
    method_info = getattr(cls, _METHOD_INFO_ATTR)
    if func_name:
      return method_info[func_name]
    return method_info


# Helper function to get the inference params
def get_inference_params(request) -> dict:
  """Get the inference params from the request."""
  inference_params = {}
  if request.model.model_version.id != "":
    output_info = request.model.model_version.output_info
    output_info = json_format.MessageToDict(output_info, preserving_proto_field_name=True)
    if "params" in output_info:
      inference_params = output_info["params"]
  return inference_params


class _MethodInfo:

  def __init__(self, method):
    self.name = method.__name__
    self.signature = build_function_signature(method)
    self.python_param_types = {
        p.name: p.annotation
        for p in inspect.signature(method).parameters.values()
        if p.annotation != inspect.Parameter.empty
    }
    self.python_param_types.pop('self', None)
