# Copyright (c) 2022-2024 Broadcom. All Rights Reserved.
# Broadcom Confidential. The term "Broadcom" refers to Broadcom Inc.
# and/or its subsidiaries.

"""Functionalities to retrieve DPU info and attributes.
"""

from esxutils import EsxcliError, runCli

from ..Errors import DpuInfoError

import logging

log = logging.getLogger(__name__)

ALIAS = 'Device Alias'
ID = 'Id'
LOCAL_EP_ADDR = 'Local Endpoint Address'
MANAGEMENT_STATE = 'Management State'
OS_VERSION = 'OS Version'
TRANSPORT_ADDR = 'Transport Address'
MANAGED = 'Managed'
UNAVAILABLE = 'Unavailable'
DEFAULT_OS_VERSION = 'N/A'

def getDpuInfo():
   """Returns a list of DPU info dicts.

      Dictionary keys which are of importance -
      "Management State" - Current state of ESXio with respect to ESXi.
                           Depends on comm-channel and trust setup status
                           between ESXio and ESXi.

      "Managed" - True if comm-channel was UP and trust setup was successful
                  between ESXi and ESXio at least once. False otherwise.
                  When True, it means that the DPU is managed by the host i.e.,
                  this server is a Monterey server for all purposes.

                  The value of this field *may* not reflect the current runtime
                  state. Use "Management State" as appropriate.

      "OS Version" - Similar to "Managed" key, it reflects a value that was
                     seen by either side in the heartbeat messages of comm
                     channel. The value of this field *may* not reflect the
                     current runtime state.
      Example:
         [
            {
               'Id': 'dpu0',
               'Device Alias': 'vmdpu0'
               'Model': 'NVIDIA-mock',
               'Vendor': 'VMware'
               'Base PCI Address': '0000:0b:00.0',
               'Firmware Version': '5.1.3.67890',
               'OS Version': 'VMware ESXio 7.0.3'/'(N/A)'
               'Transport Address': '10.185.17.1',
               'Local Endpoint Address': '10.184.108.241',
               'Address Family': 'IPv4',
               'Netstack Instance': 'defaultTcpipStack',
               'Management State': 'Managed',
               'Primary': True/False,
               'Managed': True/False,
               'API State': 'Available/Unavailable',
               'DVS Config Id": '420d244c-5c1d-ed11-8000-e8ebd3fb99d6'
            }
         ]
   """
   try:
      dpuList = runCli(['hardwareinternal', 'dpu', 'list'], True)
      log.debug('Found %d DPUs, DPU info: %s', len(dpuList), str(dpuList))
      return dpuList
   except EsxcliError as e:
      raise DpuInfoError('Failed to get DPU list: %s' % str(e))

def getManagedDpuInfo(skipUnavailableDpu=False):
   """Return a list of managed DPU info dicts.
      See also getDpuInfo().
   """
   try:
      dpuInfo = getDpuInfo()
      result = list()
      for dpu in dpuInfo:
         log.info("DPU Id: %s, State: %s, Local Endpoint Address: %s, "
                  "Transport Address: %s", dpu[ALIAS], dpu[MANAGEMENT_STATE],
                  dpu[LOCAL_EP_ADDR], dpu[TRANSPORT_ADDR])
         if dpu[MANAGEMENT_STATE] == MANAGED:
           result.append(dpu)
         elif dpu[MANAGEMENT_STATE] == UNAVAILABLE:
            # Consider only the DPUs which are managed by the host.
            if MANAGED in dpu:
               # Use the 'managed' flag that is present 8.0 U2 and later.
               if dpu[MANAGED]:
                  if not skipUnavailableDpu:
                     result.append(dpu)
                  else:
                     log.warning('Skipping DPU %s as it is UNAVAILABLE.',
                                 dpu[ALIAS])
               else:
                  log.warning('DPU %s is not a managed DPU and is ignored.',
                              dpu[ALIAS])
            elif dpu[OS_VERSION] != DEFAULT_OS_VERSION:
               # For earlier 8.0, use the OS version which is N/A for unmanaged
               # DPU.
               if not skipUnavailableDpu:
                  result.append(dpu)
               else:
                  log.warning('Skipping UNAVAILABLE DPU %s that has OS '
                              'version %s.', dpu[ALIAS], dpu[OS_VERSION])
            else:
               log.warning('Unavailable DPU %s without OS version is ignored.',
                           dpu[ALIAS])
         else:
            log.warning('DPU %s has a management state of %s and is ignored.',
                        dpu[ALIAS], dpu[MANAGEMENT_STATE])
      return result
   except KeyError as e:
      raise DpuInfoError('Attribute %s not found' % str(e))

def getDpuID(dpuInfo):
   """ Helper function to get the DPU ID.
   """
   try:
      return dpuInfo[ID]
   except KeyError as e:
      msg = 'Failed to get ID from dpu info'
      raise DpuInfoError(msg, str(dpuInfo))

def getDpuTransportAddr(dpuInfo):
   """ Helper function to get DPU transport address.
   """
   try:
      return dpuInfo[TRANSPORT_ADDR]
   except KeyError as e:
      msg = 'Failed to get transport addr from dpu info'
      raise DpuInfoError(msg, str(dpuInfo))

def getDpuAlias(dpuInfo):
   """ Helper function to get Device Alias.
   """
   try:
      return dpuInfo[ALIAS]
   except KeyError as e:
      msg = 'Failed to get Device Alias from dpu info'
      raise DpuInfoError(msg, str(dpuInfo))

def getDpuInfoFromId(dpuList, dpuId):
   """ Helper function to retrieve Dpu Info from dpu id.
   """
   for dpu in dpuList:
      if dpuId == getDpuID(dpu):
         return dpu
   raise DpuInfoError('Failed to get DPU info for DPU id: %s' % str(dpuId))

def isDpuUnavailable(dpuInfo):
   """ Whether the dpu is in UNAVAILABLE state or not.
   """
   try:
      return dpuInfo[MANAGEMENT_STATE] == UNAVAILABLE
   except KeyError as e:
      return False
