# Copyright 2019 VMware, Inc.
# All rights reserved.  -- VMware Confidential

"""Helper functions on VSAN disks.
"""
import os

from esxutils import runCli, EsxcliError

def isVsanPresent():
   """Return True if the VSAN kernel module is present.
   """
   VSAN_MOD_PATH = '/usr/lib/vmware/vmkmod/vsan'
   return os.path.isfile(VSAN_MOD_PATH)


class VsanDisk(object):
   """Class to represent a VSAN disk.
   """

   def __init__(self, diskName, diskInfo=None):
      if diskInfo is None:
         try:
            self._info = runCli(['vsan', 'storage', 'list', '-d', diskName],
                                evalOutput=True)[0]
         except EsxcliError as e:
            # Raise ValueError to distinguish a disk that is not in vSAN.
            if 'Unable to find device with name' in str(e):
               raise ValueError('%s: not a VSAN disk' % diskName)
            else:
               raise
      else:
         self._info = diskInfo

   @property
   def name(self):
      """Device name (e.g. 'mpx.vmhba2:C0:T1:L0').
      """
      return self._info['Device']

   @property
   def diskGroupUuid(self):
      """VSAN disk group UUID.
      """
      return self._info['VSAN Disk Group UUID']

   @property
   def isCapacityTier(self):
      """True if the disk is a 2nd-tier (a.k.a. capacity tier) VSAN disk.
      """
      return self._info["Is Capacity Tier"]


def getVsanDisks():
   """Get VSAN disks as a dict indexed by disk names.
   """
   disks = runCli(['vsan', 'storage', 'list'], evalOutput=True)
   return {d['Device']: d for d in disks}

def reclaimVsanDisk(diskName):
   """Reclaim (remove) a disk from a VSAN disk group.

   If the disk is first-tier, all backing 2nd-tier disks are also removed.
   If the disk group size is exactly 2, all disks are removed in the disk group,
   """
   diskGroupSize = 0
   firstTierDisk = None

   vsanDisks = getVsanDisks()
   try:
      diskInfo = vsanDisks[diskName]
      disk = VsanDisk(diskName, diskInfo=diskInfo)
   except KeyError:
      raise ValueError("%s: not a VSAN disk" % diskName)

   for name, diskInfo in vsanDisks.items():
      vsanDisk = VsanDisk(name, diskInfo=diskInfo)
      if vsanDisk.diskGroupUuid == disk.diskGroupUuid:
         diskGroupSize += 1

         if not vsanDisk.isCapacityTier:
            assert firstTierDisk is None, ("VSAN disk group %s has more "
                                           "than one first-tier SSD" %
                                           disk.diskGroupUuid)
            firstTierDisk = vsanDisk

   assert firstTierDisk is not None, ("no first-tier SSD found for VSAN disk "
                                      "group %s" % disk.diskGroupUuid)

   # Mount the group first.
   runCli(['vsan', 'storage', 'diskgroup', 'mount', '-s', firstTierDisk.name])

   if diskGroupSize == 2 or not disk.isCapacityTier:
      runCli(['vsan', 'storage', 'remove', '-s', firstTierDisk.name])
   else:
      runCli(['vsan', 'storage', 'remove', '-d', disk.name])
