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

"""Functions related to scratch on ESXi host.
"""
import os

from systemStorage import FS_TYPE_VMFS, OSDATA_LINK
from systemStorage.esxdisk import EsxDisk
from systemStorage.esxfs import getFssVolumes
import libconfigstorepy as store
import featureState
import logging

LOCKER_DOT_CONF = "/etc/vmware/locker.conf"
log = logging.getLogger(os.path.basename(__file__))

SCRATCH_SCHEMA_KEY =  "scratch_location"
SCRATCH_LOCATION = "path"
SCRATCH_CONFIG_GROUP = "storage"
SCRATCH_CONFIG_COMPONENT = "esx"
def getDefaultScratchPath():
   """Retrieve the default scratch path.

   In a booted system, the OSdata symlink points to the default scratch
   location, which can be a VMFS-L partition or a datastore .locker path.
   OSdata path is not affected by the scratch location advanced config.
   """
   if os.path.islink(OSDATA_LINK):
      return os.path.realpath(OSDATA_LINK)
   raise OSError("OSdata link does not exist")

def setScratchDir(path, context = "external"):
   """Configure the path to the scratch volume.

   Passing None to this function allows to clear (reset) the path to /scratch.
   Reboot is required for the new path to take effect.
   """
   if featureState.getVmkapps_ConfigManagerScratchConfig():
      Store = store.ConfigStore.GetStore('')
      try:
         csoId= store.ConfigStoreObjectId(SCRATCH_SCHEMA_KEY,
               SCRATCH_CONFIG_GROUP, SCRATCH_CONFIG_COMPONENT)
         cso = store.ConfigStoreObject(csoId)
         cso.SetValue(SCRATCH_LOCATION, os.path.realpath(path))
         if context == "internal":
            Store.Set(cso, store.USERDATA_CONTEXT.INTERNAL)
         else:
            Store.Set(cso)
      except Exception as exc:
         log.exception('Failed to update scratch location  with exception: %s',
                       str(exc))
   else:
      if path is None:
         open(LOCKER_DOT_CONF, 'w').close()
      else:
         os.makedirs(path, exist_ok=True)
         with open(LOCKER_DOT_CONF, 'w') as lockerConf:
            lockerConf.write("%s 0\n" % os.path.realpath(path))

def getConfiguredScratchDir(booting = False):
   """Return the path to scratch as configured in locker.conf.
   """
   if featureState.getVmkapps_ConfigManagerScratchConfig():
      Store = store.ConfigStore.GetStore('')
      try:
         csoId= store.ConfigStoreObjectId(SCRATCH_SCHEMA_KEY,
               SCRATCH_CONFIG_GROUP, SCRATCH_CONFIG_COMPONENT)
         if booting:
            configObject = store.GetDesiredUserConfig(csoId)
         else:
            configObject = Store.Get(csoId)
         if configObject is None:
            return ""
         try:
            store.SetBootStatus(csoId, True)
         except Exception as exc:
            log.exception('Failed to set boot status  with exception: %s',
                          str(exc))
         return configObject.GetValue(SCRATCH_LOCATION)
      except Exception as exc:
         log.exception('Failed to read scratch location  with exception: %s',
                       str(exc))
         return ""
   else:
      try:
         with open(LOCKER_DOT_CONF, 'r') as lockerConf:
            data = lockerConf.read()
         return data.strip().partition(' ')[0]
      except Exception:
         return ""

def getDatastoreScratchPath():
   """Identify a datastore path that can be used for scratch.

   This is to support the legacy behavior of scratch assignment when there's no
   high quality persistent storage other than datastores.
   """
   MIN_SCRATCH_DS_SIZE = 4000000000

   vols = sorted(getFssVolumes(fsType=FS_TYPE_VMFS, localOnly=True),
                 key=lambda v: v.diskName)
   for vol in vols:
      if vol.size < MIN_SCRATCH_DS_SIZE:
         continue

      # verify disk is not USB and is local (KB-1033696)
      disk = EsxDisk(vol.diskName, devInfoOnly=True)
      if not disk.isUsb and disk.isLocal:
         return os.path.join(vol.path, ".locker")

   return ""
