import yaml
import json
import numpy as np
import os

# triton_python_backend_utils is available in every Triton Python model. You
# need to use this module to create inference requests and responses. It also
# contains some utility functions for extracting information from model_config
# and converting Triton input/output types to numpy types.

import triton_python_backend_utils as pb_utils
from functools import lru_cache

def get_nested_yaml_value(file_path, keys):
    """
    从指定的YAML文件中读取并返回嵌套键对应的值。

    参数:
    file_path (str): YAML 文件的路径。
    keys (list of str): 要查找的嵌套键序列。

    返回:
    value: 对应于给定嵌套键的值。
    """
    with open(file_path, 'r', encoding='utf-8') as file:
        data = yaml.safe_load(file)
        for key in keys:
            if isinstance(data, dict) and key in data:
                data = data[key]
            else:
                return None
        return data

def get_model_info_list(model_repository_dir, model_capability):
    vendor_data = {
    "vendor": "",
    "models": []
    }
    model_dir_path = os.path.dirname(model_repository_dir)
    # Iterate over all files in the model_conf directory
    for model_dir in os.listdir(model_dir_path):
        file_path = os.path.join(model_dir_path, model_dir, "model_description.yaml")
        if os.path.exists(file_path):
            with open(file_path, 'r', encoding='utf-8') as file:
                data = yaml.safe_load(file)
                # Check if required keys exist and match the filter
                if isinstance(data, dict) and all(key in data for key in ('model_name', 'model_capability', 'model_is3rdModel')):
                    if model_capability in data['model_capability']:
                        model_name = data.get('model_name')
                        display_name = data.get('model_display_name')
                        # 优先使用 model_display_name，否则使用 model_name 的后半部分
                        name_value = display_name or "_".join(model_name.split("_", 1)[1:]) if model_name else ""
                        vendor_data["models"].append({
                            "model": model_name,
                            "name": name_value,
                            "capability": data.get('model_capability'),
                            "subConfig": [],
                            "debName": data.get('model_deb_name'),
                            "debVersion": "",
                            "is3rdModel": data.get('model_is3rdModel')
                        })

    json_response = json.dumps(vendor_data, ensure_ascii=False)
    return json_response

@lru_cache(maxsize=None)
def model_bank_process(model_repository, config_file_path, input_key_strings, input_engine_strings):
  if input_key_strings == "model_info_list":
    model_list = get_model_info_list(model_repository, input_engine_strings)
    value_bytes = np.array([model_list.encode('utf-8')], dtype=object)  # Use dtype=object for BYTES type
    out_tensor = pb_utils.Tensor("value", value_bytes)
    conf_response = pb_utils.InferenceResponse(output_tensors=[out_tensor])
  else:
    with open(config_file_path, 'r', encoding='utf-8') as file:
      data = yaml.safe_load(file)
      if input_key_strings in data:
        value = data[input_key_strings]
      else:
        value = ""

    print("value: ", value)

    value_bytes = np.array([value.encode('utf-8')], dtype=object)  # Use dtype=object for BYTES type

    out_tensor = pb_utils.Tensor("value", value_bytes)

    conf_response = pb_utils.InferenceResponse(output_tensors=[out_tensor])
  return conf_response

class TritonPythonModel:
    """Your Python model must use the same class name. Every Python model
    that is created must have "TritonPythonModel" as the class name.
    """

    def initialize(self, args):
        """`initialize` is called only once when the model is being loaded.
        Implementing `initialize` function is optional. This function allows
        the model to intialize any state associated with this model.
        Parameters
        ----------
        args : dict
          Both keys and values are strings. The dictionary keys and values are:
          * model_config: A JSON string containing the model configuration
          * model_instance_kind: A string containing model instance kind
          * model_instance_device_id: A string containing model instance device ID
          * model_repository: Model repository path
          * model_version: Model version
          * model_name: Model name
        """
        self.model_config = json.loads(args["model_config"])
        self.model_repository = args["model_repository"]

    def execute(self, requests):
        """`execute` MUST be implemented in every Python model. `execute`
        function receives a list of pb_utils.InferenceRequest as the only
        argument. This function is called when an inference request is made
        for this model. Depending on the batching configuration (e.g. Dynamic
        Batching) used, `requests` may contain multiple requests. Every
        Python model, must create one pb_utils.InferenceResponse for every
        pb_utils.InferenceRequest in `requests`. If there is an error, you can
        set the error argument when creating a pb_utils.InferenceResponse
        Parameters
        ----------
        requests : list
          A list of pb_utils.InferenceRequest
        Returns
        -------
        list
          A list of pb_utils.InferenceResponse. The length of this list must
          be the same as `requests`
        """

        responses = []

        config_dir = os.path.join(self.model_repository, '1')
        config_file_path = os.path.join(config_dir, 'default_model.yaml')
        if os.path.exists(config_file_path):
          with open(config_file_path, 'r', encoding='utf-8') as file:

            data = yaml.safe_load(file)
            if isinstance(data, dict):
            # Every Python backend must iterate over everyone of the requests
            # and create a pb_utils.InferenceResponse for each of them.
              for request in requests:
                  # Get INPUT0
                  input_key_strings = pb_utils.get_input_tensor_by_name(request, "key").as_numpy()
                  input_key_strings = input_key_strings[0].decode('utf-8')
                  input_engine_strings = pb_utils.get_input_tensor_by_name(request, "engine").as_numpy()
                  input_engine_strings = input_engine_strings[0].decode('utf-8')
                  print(f"input_key_strings:{input_key_strings}")
                  print(f"input_engine_strings:{input_engine_strings}")
                  conf_response = model_bank_process(self.model_repository, config_file_path, input_key_strings,input_engine_strings)
                  responses.append(conf_response)

        # You should return a list of pb_utils.InferenceResponse. Length
        # of this list must match the length of `requests` list.
        return responses

    def finalize(self):
        """`finalize` is called only once when the model is being unloaded.
        Implementing `finalize` function is OPTIONAL. This function allows
        the model to perform any necessary clean ups before exit.
        """
        print('Cleaning up...')
