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

"""Utility module for ESX installation.
"""
import logging
import os
import shutil

from coredump import unconfigureCoredumpToFile, unconfigureCoredumpToPart
from vmSyslogUtils import (setSyslogLogDir, resetSyslogLogDir, startSyslog,
                           stopSyslog)

from systemStorage import (BOOTBANK_LINK, ALTBOOTBANK_LINK, SCRATCH_LINK,
                           STORE_LINK, VARLOG_LINK, FS_TYPE_VFAT,
                           FS_TYPE_FAT16, FS_TYPE_FAT32)
from systemStorage.autoPartition import autoPartition, partitionOsdataDisk
from systemStorage.esxfs import fsRescan, umountFileSystems
from systemStorage.upgradeUtils import copyFiles
from systemStorage.vfat import mcopy

log = logging.getLogger(os.path.basename(__file__))
log.setLevel(logging.DEBUG)

def decommissionSystemDisk():
   """Decommission the system disk set up by system-storage jumpstart.

   This prepares the environment for installer to perform disk partitioning.

   Actions:
    - Deconfigure coredump to partition and file.
    - Set /bootbank to /tmp, /store and /scratch to /tmp/scratch.
    - Restart syslog to log to /scratch/log (/tmp/scratch/log).
    - Remove /altbootbank symlink.
   """
   unconfigureCoredumpToFile()
   unconfigureCoredumpToPart()

   tmpDir = os.path.join(os.path.sep, 'tmp')
   scratchDir = os.path.join(tmpDir, 'scratch')
   logDir = os.path.join(scratchDir, 'log')
   os.makedirs(logDir, exist_ok=True)

   if os.path.realpath(SCRATCH_LINK) != scratchDir:
      stopSyslog()
      # Copy the current log in /var/log.
      curLogDir = VARLOG_LINK
      copyFiles(curLogDir, logDir)
      # Reset current scratch.
      if os.path.islink(SCRATCH_LINK):
         os.unlink(SCRATCH_LINK)
      elif os.path.isdir(SCRATCH_LINK):
         # XXX: system-storage jumpstart does not handle legacy layout,
         # /scratch will be a folder created by logging.
         shutil.rmtree(SCRATCH_LINK)
      os.symlink(scratchDir, SCRATCH_LINK)
      # Restart syslog to continue from the new scratch.
      setSyslogLogDir(logDir, reloadService=False)
      startSyslog()
      # Reset logdir to default without reloading syslog, this clears the logdir
      # config in configstore without affecting logging in this session.
      resetSyslogLogDir(reloadService=False)

   if os.path.islink(STORE_LINK):
      os.unlink(STORE_LINK)
      os.symlink(scratchDir, STORE_LINK)

   if os.path.islink(BOOTBANK_LINK):
      os.unlink(BOOTBANK_LINK)
      os.symlink(tmpDir, BOOTBANK_LINK)

   if os.path.islink(ALTBOOTBANK_LINK):
      os.unlink(ALTBOOTBANK_LINK)

def copyFileFromFATPart(disk, srcPath, destPath):
   """mcopy a file from a FAT partition on the given disk.

   On finding the first partition that has the file, copy will be performed
   using mcopy. An exception is raised of none of the disk's FAT partitions
   has the file or if the disk does not have any FAT partitions.
   """
   # For GPT, basic data type is mapped to VFAT.
   fatFsTypes = (FS_TYPE_VFAT, FS_TYPE_FAT16, FS_TYPE_FAT32)

   disk.scanPartitions()
   for part in disk.getPartitionsByFsTypes(fatFsTypes):
      offset = part.start * disk.sectorSize
      try:
         mcopy(disk.path, [srcPath], destPath, byteOffset=offset, srcIsPart=True)
         break
      except RuntimeError:
         pass
   else:
      raise RuntimeError('Disk %s does not contain file %s in any of its FAT '
                         'partitions' % (disk.name, srcPath))

def autopartitionBootDisk(disk, keepDatastore, createDatastore):
   """Partition the ESX boot disk.

   This function is called by Weasel to install a brand new partition table on
   the ESX boot disk.
   """
   umountFileSystems(disk.name)
   log.info('%s: partitioning boot disk (keepDatastore=%s, createDatastore=%s)',
            disk.name, keepDatastore, createDatastore)
   autoPartition(disk, keepDatastore=keepDatastore,
                 createDatastore=createDatastore)
   fsRescan()

def unmountAndPartitionOsdataDisk(disk):
   """Partition the disk selected for ESX OSData.

   This function is called by Weasel to install a brand new partition table on
   the disk selected for ESX OSData in system with USB/SD boot device.
   """
   umountFileSystems(disk.name)
   log.info('%s: partitioning ESX OSData disk', disk.name)
   disk.clearPartitions()
   partitionOsdataDisk(disk, doPartSync=True)
