from clarifai_grpc.grpc.api import resources_pb2, service_pb2
from collections.abc import Iterator
from google.protobuf import json_format

from clarifai.client.runner import Runner


class MyRunner(Runner):
  """A custom runner that adds "Hello World" to the end of the text and replaces the domain of the
  image URL as an example.
  """

  def run_input(self, input: resources_pb2.Input, output_info: resources_pb2.OutputInfo,
                **kwargs) -> resources_pb2.Output:
    """This is the method that will be called when the runner is run. It takes in an input and
    returns an output.
    """

    output = resources_pb2.Output()

    data = input.data

    # Optional use of output_info
    params_dict = {}
    if "params" in output_info:
      params_dict = output_info["params"]

    if data.text.raw != "":
      output.data.text.raw = data.text.raw + "Hello World" + params_dict.get(
          "hello", "") + kwargs.get("extra", "")
    if data.image.url != "":
      output.data.text.raw = data.image.url.replace("samples.clarifai.com",
                                                    "newdomain.com" + params_dict.get("domain",))
    return output

  def generate(self, request: service_pb2.PostModelOutputsRequest
              ) -> Iterator[service_pb2.MultiOutputResponse]:
    """Example yielding a whole batch of streamed stuff back.
    """

    model = request.model
    output_info = None
    if request.model.model_version.id != "":
      output_info = json_format.MessageToDict(
          model.model_version.output_info, preserving_proto_field_name=True)

    for i in range(10):  # fake something iterating generating 10 times.

      outputs = []
      for input in request.inputs:
        # output = self.run_input(input, output_info, extra=f" {i}")
        output = resources_pb2.Output()
        output.data.text.raw = f"Generate Hello World {i}"
        outputs.append(output)
      resp = service_pb2.MultiOutputResponse(outputs=outputs,)
      yield resp

  def stream(self, request: service_pb2.PostModelOutputsRequest
            ) -> Iterator[service_pb2.MultiOutputResponse]:
    """Example yielding a whole batch of streamed stuff back.
    """

    model = request.model
    output_info = None
    if request.model.model_version.id != "":
      output_info = json_format.MessageToDict(
          model.model_version.output_info, preserving_proto_field_name=True)

    for i in range(10):  # fake something iterating generating 10 times.

      outputs = []
      for input in request.inputs:
        # output = self.run_input(input, output_info, extra=f" {i}")
        output = resources_pb2.Output()
        out_text = input.data.text.raw + f"Stream Hello World {i}"
        print(out_text)
        output.data.text.raw = out_text
        outputs.append(output)
      resp = service_pb2.MultiOutputResponse(outputs=outputs,)
      yield resp


if __name__ == '__main__':
  # Make sure you set these env vars before running the example.
  # CLARIFAI_PAT
  # CLARIFAI_USER_ID

  # You need to first create a runner in the Clarifai API and then use the ID here.
  MyRunner(runner_id="matt-test-runner", base_url="http://q6:32013", num_parallel_polls=1).start()
