#########################################################################
#
#   Copyright 2016-2022 VMware, Inc.  All rights reserved. VMware Confidential
#
#########################################################################
from extensions import extend, Hook
from vcsa_utils import getComponentDiscoveryResult, isDisruptiveUpgrade
from fss_utils import getTargetFSS

from importlib import import_module
from logging import getLogger
from os import environ, chmod
from os.path import join, isfile
from sys import path as PYTHONPATH
from shutil import copyfile
PYTHONPATH.append(environ['VMWARE_PYTHON_PATH'])
from .patches.patchUtil import (persistCurrentVersionFile, getSourceVersion,
                               getTargetVersion, TO_PREPARE_SUFFIX,
                               getLocalRuntimeModification,
                               getTargetRuntimeModification,
                               updateRuntimeModification)

# eam constsants
EAM_SERVICES_JSON_NAME = "eam"

# patch constants
NDU_FSS = "NDU_Limited_Downtime"
PATHCES_MOD = "patches"
PATCH_PREFIX = "patch"
PATCH_EXTENSION_RUNTIME_MOD = "PATCH_EXTENSION"

_log = getLogger(__name__)

# TODO: remove usage of isExpand once NDU_Limited_downtime is enabled

@extend(Hook.Discovery)
def discover(ctx):
   '''
   Discovers ESX Agent Manager previous installed
   version and decides whether to patch the configuration
   or not and how to do it.
   '''
   _log.info("Starting discovery for vSphere ESX Agent Manager ..")
   componentDiscoveryResult = \
      getComponentDiscoveryResult(EAM_SERVICES_JSON_NAME)

   if not getTargetFSS(NDU_FSS) or isDisruptiveUpgrade(ctx):
      if componentDiscoveryResult.patchable :
         persistCurrentVersionFile()
   else:
      repCfg = _replicationConfig(True)
      _addRuntimeModification(repCfg)
      _log.info("Created replication config {0}".format(repCfg))
      componentDiscoveryResult.replicationConfig = repCfg

   return componentDiscoveryResult

@extend(Hook.Expand)
def expand(ctx):
   '''
   Expands patches from previous ESX Agent Manager version to the current.
   @param ctx, patch context
   '''
   if getTargetFSS(NDU_FSS):
      _log.info("Starting Expand phase for vSphere ESX Agent Manager ..")
      _replicateExpandConfigFiles(True)
      _doPatch(ctx, True)
      updateRuntimeModification((PATCH_EXTENSION_RUNTIME_MOD,))
      _log.info(
         "Finish with Expand phase. vSphere ESX Agent Manager Expand done."
      )

@extend(Hook.Patch)
def patch(ctx):
   '''
   Patches from the previous ESX Agent Manager version to the current.
   @param ctx, patch context
   '''
   _log.info("Starting Patch phase for vSphere ESX Agent Manager ..")
   if not getTargetFSS(NDU_FSS):
      _doPatch(ctx, False)
   elif isDisruptiveUpgrade(ctx):
      _replicateExpandConfigFiles(False)
   else:
      _log.info("Nothing to do in Patch phase when NDU.")
      #do nothing
   _log.info("Finish with patch phase. vSphere ESX Agent Manager Patch done.")

def _replicateExpandConfigFiles(isExpand):
   repCfg = _replicationConfig(isExpand)
   _log.info("Created replication config {0}".format(repCfg))
   if not isExpand:
      _addRuntimeModification(repCfg)

   expandFiles = filter(lambda k: k.endswith(TO_PREPARE_SUFFIX), repCfg.keys())
   for ef in expandFiles:
      _log.info("Expand file: {0}".format(ef))
      if isExpand:
         if isfile(repCfg[ef]):
            _log.info("Expand Copy {0} to {1}".format(repCfg[ef], ef))
            copyfile(repCfg[ef], ef)
      else:
         _log.info("Patch Copy {0} to {1}".format(ef, repCfg[ef]))
         copyfile(ef, repCfg[ef])
         if ef == getLocalRuntimeModification():
            chmod(repCfg[ef], 0o644)

def _replicationConfig(isExpand):
   replicationConfig = {}
   _forEachApplicablePatch(
      lambda patch: patch.updateReplicationConfig(replicationConfig),
      isExpand,
      "replicationConfig"
   )
   return replicationConfig

def _addRuntimeModification(repCfg):
   repCfg[getLocalRuntimeModification()] = getTargetRuntimeModification()

def _doPatch(ctx, isExpand):
   if not isExpand:
      _updateEamExtension(ctx)
   _log.info("Starting EAM patching..")
   _forEachApplicablePatch(
      lambda patch: patch.run(ctx, isExpand),
      isExpand,
      "run"
   )

def _forEachApplicablePatch(applyFn, isExpand, opMsg):
   # iterate over patches and pick ones which are not currently applied
   for i in range(getSourceVersion(isExpand), getTargetVersion()):
      patchMod = PATCH_PREFIX + str(i + 1)
      try:
         # apply patch for given version
         applyFn(_getModule(patchMod))
         _log.info("Successfully applied {0} on {1}".format(opMsg,patchMod))
      except Exception as e:
         _log.exception("Applying {0} on {1} failed".format(opMsg, patchMod))
         raise e

def _updateEamExtension(ctx):
   # Special patch for updating EAM extension.
   patchMod = _getModule('patchExtension')
   patchMod.run(ctx)

def _getModule(patchName):
   return import_module('.' + patchName, __name__ + '.' + PATHCES_MOD)
