#
# Copyright (c) 2021 Nutanix Inc. All rights reserved.
#
# Author: reuben.cherian@nutanix.com
#
# Helper Tools for NVMe Firmware using IntelMAS Utility
#

import json
import os
import tarfile
import nutanix.tools as tools
import nutanix.hypervisor_tools as hyp_tools

from distutils.version import StrictVersion


HYP_TOOL_MAP = {
"ahv": "/opt/intel/nvmupdate/nvmupdate64e",
"esx": "/usr/bin/nvmupdaten64e"
}


class NvmTools(object):
  def __init__(self, extract=True):
    self.nvm_path = os.path.dirname(os.path.abspath(__file__))
    self.extract = extract
    self.nvm_util = ""
    self.nic_cmds = {
        "INVENTORY": "%s -i -l"
    }

  def get_all_nics(self, supported_nics):
    """
    Fetches all NIC details using nvm tool.
    args: supported_nics is a list of models.
          ['X710','X550']
    Returns:
       - Dictionary of NIC details on success:
      (True, filtered_nics )
       Ex: {'X550': [{'CIVD':'1.2778.0','fw_version':'2.20'}],
            'X710': [{'CIVD':'1.3126.0','fw_version':'7.00'}]}
       - Error on Failure:(False, err)
     nvmupdate64e Output:
     ./nvmupdate64e -i -l
     Intel(R) Ethernet NVM Update Tool
     NVMUpdate version 1.35.23.3
     Copyright (C) 2013 - 2020 Intel Corporation.
     [00:024:00:00]: Intel(R) Ethernet Converged Network Adapter X550-T2
        Vendor                 : 8086
        Device                 : 1563
        Subvendor              : 8086
        Subdevice              : 001A
        Revision               : 1
        LAN MAC                : A0369FF30364
        Alt MAC                : A0369FF30364
        SAN MAC                : A0369FF30366
        ETrackId               : 8000113B
        NVM Version            : 2.32(2.20)
        PBA                    : H49289-010
        VPD status             : Valid
        VPD size               : 38
        NVM update             : No config file entry
          checksum             : Valid
        OROM update            : No config file entry
          CIVD                 : 1.2684.0
    """
    ret, env, hypervisor = hyp_tools.get_hypervisor_tools()
    if not ret:
      tools.module_print("Failed to get hypervisor: %s" % env)
      return (False, env)

    nvmupdate_tool = HYP_TOOL_MAP.get(env)
    if os.path.exists(nvmupdate_tool):
      tools.module_print("nvmupdate tool exists at: %s" % nvmupdate_tool)
      self.nvm_util = nvmupdate_tool
    else:
      # Extract the tool
      tools.module_print("%s path doesn't exists. So Proceeding for Extraction" % nvmupdate_tool)
      rc, path = self._extract_tool()
      if not rc:
        return (False, "Exception in extracting nvmupdate bundle")
      self.nvm_util = path

    cmd = self.nic_cmds.get("INVENTORY") % self.nvm_util
    tools.module_print("Running cmd: %s" % cmd)
    out, err, rc = tools.run_cmd(cmd, no_error=True)
    if rc:
      tools.module_print("Could not execute command: %s. Error: %s" %
                         (cmd, err))
      return (False, err)
    try:
      nic_dict = self._get_dict_from_out(out)
    except Exception as e:
      tools.module_print("Exception while getting NIC dict. Error: %s" % str(e))
      return (False, str(e))
    tools.module_print("NIC Dict: %s" % nic_dict)
    ret, filtered_nics = self._filter_supported_nics(nic_dict, supported_nics)
    if not ret:
      return (False, filtered_nics)
    return (True, filtered_nics)

  def _extract_tool(self):
    """
    Untars the nvmupdate utility bundle and returns path.
    Returns:
      (True, path) on success.
      (False, "") on failure.
    """
    manifest_path = os.path.join(self.nvm_path, "manifest.json")
    if not os.path.exists(manifest_path):
      tools.module_print("nvmupdate utility manifest does not exist at %s" %
                          (manifest_path))
      return (False, "")
    tools.module_print("Found the nvmupdate manifest %s" % manifest_path)
    try:
      with open(manifest_path, 'r') as f:
        image_metadata = json.load(f)
        binary_file = image_metadata.get("file_name")
        nvm_tar_bundle = os.path.join(tools.get_root(), binary_file)

      if not os.path.exists(nvm_tar_bundle):
        tools.module_print("Nvmupdate tar does not exist at %s" %
                            nvm_tar_bundle)
        return (False, "")

      tools.module_print("Found the nvmupdate tar %s" % nvm_tar_bundle)
      ret, env, hypervisor = hyp_tools.get_hypervisor_tools()
      if not ret:
        tools.module_print("Failed to get hypervisor: %s" % env)
        return (False, env)

      binary = os.path.join(env, HYP_TOOL_MAP[env].split('/')[-1])
      nvm_util = os.path.join(self.nvm_path, binary)
      if os.path.exists(nvm_util):
        tools.module_print("nvmupdate utility already exists %s" %
                            nvm_util)
        return (True, nvm_util)
      with tarfile.open(nvm_tar_bundle) as tar:
        tools.module_print("Extracting nvmupdate utility bundle")
        tar.extractall(path=self.nvm_path)
      if not os.path.exists(nvm_util):
        tools.module_print("nvmupdate utility does not exist at %s" %
                            nvm_util)
        return (False, "")
      tools.module_print("Found the nvmupdate utility %s" % nvm_util)
    except Exception as e:
      tools.module_print("Exception while extracting nvmupdate [%s]" % str(e))
      return (False, "")
    return (True, nvm_util)

  def _get_dict_from_out(self, out, out_dict={}, sep=":"):
    """
    Returns the list of supported NIC details.
    """
    key = ""
    main_dict = {}
    out_dict = {}
    if isinstance(out, list):
      lines = out
    else:
      lines = out.splitlines()
    for line in lines:
      if line.startswith('['):
        key = line.split(':')[-1].strip()
        out_dict = {}
        main_dict[key] = out_dict
        continue
      if sep in line:
        key, value = line.split(sep, 1)
        if key.strip() in out_dict:
          out_dict[key.strip()] += value.strip() + ","
        elif key:
          out_dict[key.strip()] = value.strip()
    return main_dict

  def _filter_supported_nics(self, nic_dict, supported_nics):
   """
   Filters only the details of the Fujitsu supported NICs
   returns:
      (True, filtered_dict) on success
         - {'X550':[{'CIVD':'1.2778.0','fw_version':'2.20'}]}
      (False, err) on Failure
   """
   filtered_nics = {}
   try:
     for model in nic_dict.keys():
       model_list = list(set([x for x in supported_nics if x in model]))
       if model_list:
         civd = nic_dict[model].get('CIVD')
         fw_version = nic_dict[model].get('NVM Version') \
           .split('(')[-1].strip(')')
         if model_list[0] not in filtered_nics:
           filtered_nics[model_list[0]] = []
           info = {"CIVD": civd, "fw_version": fw_version}
           filtered_nics[model_list[0]].append(info)
         else:
           if StrictVersion(fw_version) < \
             StrictVersion(filtered_nics[model_list[0]][0].get("fw_version")):
             filtered_nics[model_list[0]] = []
             info = {"CIVD": civd, "fw_version": fw_version}
             filtered_nics[model_list[0]].append(info)
   except Exception as ex:
     tools.module_print("Unable to Filter supported NICs: %s" % str(ex))
     return (False, str(ex))
   return (True, filtered_nics)
