from pathlib import Path
import subprocess
import yaml
import json
import numpy as np
import shutil
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

@lru_cache(maxsize=1024)  # 根据实际情况调整缓存大小
def get_deb_package_info(file_path):
    try:
        # 使用dpkg -S命令查找哪个已安装的.deb包拥有指定文件
        result = subprocess.run(['dpkg', '-S', file_path], 
                               stdout=subprocess.PIPE, 
                               stderr=subprocess.PIPE, 
                               text=True, 
                               check=True)

        # 解析输出，获取包名和版本号
        output = result.stdout.strip()
        if not output:
            return None, None
        
        # dpkg -S 输出格式：package_name:file_path
        package_info = output.split(':', 1)[0].strip()

        # 使用dpkg-query命令获取包的版本号
        version_result = subprocess.run(['dpkg-query', '-W', '--showformat=${Version}', package_info],
                                        stdout=subprocess.PIPE,
                                        stderr=subprocess.PIPE,
                                        text=True,
                                        check=True)
        
        version = version_result.stdout.strip()
        return package_info, version

    except subprocess.CalledProcessError as e:
        # 如果找不到对应的.deb包或者发生其他错误
        print(f"Error occurred while trying to find the deb package for {file_path}: {e.stderr}")
        return None, None

def ensure_config_file_exists(config_dir, config_file_path, model_repository_dir):
    # user_home = os.path.expanduser('~')
    # config_dir = os.path.join(user_home, '.config')
    # config_file_path = os.path.join(config_dir, 'config_bank.yaml')
    
    if not os.path.exists(config_dir):
        os.makedirs(config_dir)
    
    if os.path.exists(config_file_path):
        return True
    else:
        source_file_path = f"{model_repository_dir}/1/config_bank.yaml"
        try:
            shutil.copy(source_file_path, config_file_path)
            return True
        except Exception as e:
            print(f"Error copying file: {e}")
            return False

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_conf_path = os.path.join(f"{model_repository_dir}/1/", 'model_conf')
    # Iterate over all files in the model_conf directory
    for file_name in os.listdir(model_conf_path):
        if file_name.endswith('.yaml'):
            file_path = os.path.join(model_conf_path, file_name)
            package_info, version = get_deb_package_info(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 data['model_capability'] == model_capability:
                        vendor_data["models"].append({
                            "name": data.get('model_name'),
                            "capability": data.get('model_capability'),
                            "subConfig": [],
                            "debName": package_info,  # Assuming these fields are not provided in the YAML files,
                            "debVersion": version,  # you may need to fill them according to your requirements.
                            "is3rdModel": data.get('model_is3rdModel')
                        })
    json_response = json.dumps(vendor_data, ensure_ascii=False)
    return json_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"]
        # print(f"self.model_config:{self.model_config}")
        # print(f"self.model_repository:{self.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_file_path = f"{self.model_repository}/1/config_bank.yaml"

        user_home = os.path.expanduser('~')
        config_dir = os.path.join(user_home, '.config')
        config_file_path = os.path.join(config_dir, 'config_bank.yaml')
        if ensure_config_file_exists(config_dir, config_file_path, self.model_repository):
          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}")
                  if input_key_strings == "model_info_list":
                    model_list = get_model_info_list(self.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)
                    # print("value_bytes: ", value_bytes)

                    conf_response = pb_utils.InferenceResponse(output_tensors=[out_tensor])
                    responses.append(conf_response)
                  elif input_key_strings.startswith("change_"):
                    print(f"startswith(change_)input_key_strings:{input_key_strings}")
                    value_bytes = np.array(["change_failed".encode('utf-8')], dtype=object)
                    change_key = input_key_strings[len("change_"):]
                    with open(config_file_path, 'r', encoding='utf-8') as file:
                        data = yaml.safe_load(file)
                        if isinstance(data, dict):
                            # Update the configuration dictionary
                            data[input_engine_strings] = change_key
                            with open(config_file_path, 'w', encoding='utf-8') as file:
                                yaml.dump(data, file, default_flow_style=False)
                                value_bytes = np.array(["change_succ".encode('utf-8')], dtype=object)
                    out_tensor = pb_utils.Tensor("value", value_bytes)
                    conf_response = pb_utils.InferenceResponse(output_tensors=[out_tensor])
                    responses.append(conf_response)
                  else:
                    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)
                    # print("value_bytes: ", value_bytes)

                    conf_response = pb_utils.InferenceResponse(output_tensors=[out_tensor])
                    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...')
