"""
Copyright 2018-2021 VMware, Inc.  All rights reserved.
-- VMware Confidential
"""

import os, sys
sys.path.append(os.environ['VMWARE_PYTHON_PATH'])

from extensions import extend, Hook
from l10n import msgMetadata as _T, localizedString as _
from fss_utils import getTargetFSS
from patch_specs import DiscoveryResult
from vcsa_utils import isDisruptiveUpgrade
from transport import getCommandExitCode
from transport.local import LocalOperationsManager
import vcsa_utils
import logging
import cis.utils
import shutil

logDir = os.path.join(os.environ["VMWARE_LOG_DIR"], 'firstboot')
cis.utils.setupLogging('vsanvcmgmt_update', logMechanism='file', logDir=logDir)
cis.utils.log = lambda msg, quiet=False: logging.info(msg)
cis.utils.log_warning = lambda msg, quiet=False: logging.warning(msg)
cis.utils.log_error = lambda msg, quiet=False: logging.error(msg)

vsanhealth_cfg_dir = '/etc/vmware-vsan-health'
cns_pgpass = '.cns_pgpass'
psql = '/opt/vmware/vpostgres/current/bin/psql'
rootDir = os.path.dirname(os.path.realpath(__file__))
_CNS_DB_EXPAND_FILE = os.path.join(rootDir, "cns_expand.sql")
fileservice_ovf_dir = '/storage/updatemgr/vsan/'


@extend(Hook.Discovery)
def discover(ctx):
   # Create marker.template file using for NDU registration.
   # Need to preserve OVF files for file service during NDU.
   # More detail can be found at
   # https://confluence.eng.vmware.com/display/VCUSOF/Configuration+upgrade#Configurationupgrade-Registrations
   if getTargetFSS("NDU_Limited_Downtime") and not isDisruptiveUpgrade(ctx):
      source_marker = os.path.join(ctx.stageDirectory, "marker.template")
      with open(source_marker, "w") as f:
         f.write('')

      target_marker = "/etc/vmware-vsan-health/marker.template"
      replicationConfig = {
         source_marker : target_marker,
         fileservice_ovf_dir : fileservice_ovf_dir
      }

      return vcsa_utils.getComponentDiscoveryResult(
            'vsan-health', deprecatedProducts=["VMware-vsan-racadm"],
            replicationConfig=replicationConfig)
   else:
      return vcsa_utils.getComponentDiscoveryResult(
         'vsan-health', deprecatedProducts=["VMware-vsan-racadm"])

@extend(Hook.Prepatch)
def prePatch(ctx):
   if not getTargetFSS("NDU_Limited_Downtime"):
      """void prePatch(PatchContext sharedCtx) throw UserError"""
      cns_pgpass_f = os.path.join(vsanhealth_cfg_dir, cns_pgpass)
      if os.path.exists(cns_pgpass_f):
         logging.info('Preserving %s to %s' % (cns_pgpass_f, ctx.stageDirectory))
         shutil.copy(cns_pgpass_f, ctx.stageDirectory)

      else:
         logging.info('No CNS pgpass file found to be preserved.')

@extend(Hook.Expand)
def expand(ctx):
   cns_pgpass_f = os.path.join(vsanhealth_cfg_dir, cns_pgpass)
   if os.path.exists(cns_pgpass_f):
      logging.info('Preserving %s to %s' % (cns_pgpass_f, ctx.stageDirectory))
      shutil.copy(cns_pgpass_f, ctx.stageDirectory)
   else:
      logging.info('No CNS pgpass file found to be preserved.')

   if getTargetFSS("NDU_Limited_Downtime") and not isDisruptiveUpgrade(ctx):
      logging.info("Starting NDU expand for vsan-health.")

      """ Run DB expand script """
      cmd = [psql, '-U',  'postgres', '-d', 'VCDB', '-f', _CNS_DB_EXPAND_FILE]
      exitCode = getCommandExitCode(LocalOperationsManager(), cmd)
      if exitCode:
         cause = _(_T('cnsdb.ndu.expand.generic.error',
                      'Failed to extend the cns database state.'))
         raise UserError(cause=cause)
   else:
      logging.info("FSS is off, all work is done in Patch hook")


def patch_cns(ctx):
   if cis.utils.is_windows():
      logging.info('Skip CNS b2b upgrade on Windows.')
      return
   logging.info('Starting CNS DB configuration b2b upgrade.')
   cns_pgpass_f = os.path.join(vsanhealth_cfg_dir, cns_pgpass)
   cns_pgpass_bak = os.path.join(ctx.stageDirectory, cns_pgpass)
   if os.path.exists(cns_pgpass_bak):
      logging.info('CNS pgpass backup file of old VC found at at: %s' %
                   cns_pgpass_bak)
      logging.info('Restoring CNS pgpass file from %s to %s' % (cns_pgpass_bak,
                                                                cns_pgpass_f))
      shutil.copyfile(cns_pgpass_bak, cns_pgpass_f)
      from CnsVcConfigure import cns_configure
      cns_configure(freshinstall=False)
   else:
      logging.info('No CNS pgpass backup file of old VC at: %s' %
                   cns_pgpass_bak)
      logging.info('Generate CNS DB configuration from scratch')
      from CnsVcConfigure import cns_configure
      cns_configure()

   logging.info("CNS DB configuration b2b upgrade complete.")

@extend(Hook.Patch)
def patch(ctx):
   if getTargetFSS("NDU_Limited_Downtime") and not isDisruptiveUpgrade(ctx):
      logging.info("This is case of NDU. Work will be done in expand hook.")
      return

   logging.info("Starting B2B Upgrade for vsan-health")
   savedArgs = sys.argv[:]
   sys.argv = sys.argv[0:1]
   try:
      from cis.defaults import def_by_os, get_cis_install_dir
      sys.path.append(os.path.join(get_cis_install_dir(),
                                   def_by_os('vmware-vpx/firstboot',
                                             'firstboot')))
      from vsanhealth_firstboot import update_service_registration
      update_service_registration()
      patch_cns(ctx)

   except:
      logging.exception('Failed to patch vSAN management service')
      raise
   finally:
      sys.argv = savedArgs
