"""
Copyright 2019-2021 VMware, Inc.  All rights reserved. -- VMware Confidential
Utility Module
"""

import os
import sys
import shutil
import logging

sys.path.append(os.environ['VMWARE_PYTHON_PATH'])
from cis.utils import readprop, run_command
from cis.defaults import get_env_var, get_cis_data_dir, \
    get_cis_log_dir, get_cis_config_dir

VTSDB_HOME = get_env_var('VMWARE_POSTGRES_BASE')
sys.path.insert(0, os.sep.join([VTSDB_HOME, 'share', 'python-modules']))
from vpostgres_cis.firstboot import GetPgVersion

COMPONENT_NAME = 'vtsdb'
vtsdb_patch_key = 'vtsdb.patch.version'

logger = logging.getLogger(__name__)


def get_config_dir():
    return '/etc/vmware-%s' % COMPONENT_NAME


# The file will be installed as part of the target deployment
# in case of a nondisruptive upgrade.
# During upgrade, we need to preserve its original content in the
# @Expand hook in order to be
# available in the target appliance.
__VERSION_FILE__ = '%s/vtsdb_version.txt' % get_config_dir()
__VERSION_FILE_NAME__ = 'vtsdb_version.txt'

def getSourceVersion(parentDir, fileName=__VERSION_FILE_NAME__):
    ''' Loads current component version from file located at @parentDir
    and named @fileName
    '''
    filePath = os.path.join(parentDir, fileName)
    if os.path.exists(filePath):
            version = readprop(filePath, vtsdb_patch_key)
    else:
        version = None
    return version

def is_patch_needed(curr_version, target_patch_version):
    """
    This function will check if the patch to target_patch_version is required
    :param curr_version:
    :param target_patch_version:
    :return:
    """
    if curr_version < target_patch_version:
        return True
    return False

def preserveVersionFile(stageDir, filePath=__VERSION_FILE__):
    """ Persists the version file in the stage directory

    @param stageDir: Directory to preserve the version file in.
    """
    if os.path.exists(filePath):
        shutil.copy(filePath, stageDir)
    else:
        logger.info("Version file doesn't exist!! This may be a fresh install of vtsdb")

def GetCurrentPGVersion():
    """
    Fetch the current postgres version from /opt/vmware/vpostgres/current/bin/psql --version
    """
    vpostgres_bin = os.path.join(get_env_var('VMWARE_POSTGRES_BIN'), 'psql')
    cmd = [vpostgres_bin, '--version']
    logger.info('Running %s' % cmd)
    (ret, out, err) = run_command(cmd)
    if ret != 0:
        logger.error('Cmd returned error code %d with error %s out %s' % (ret, err, out))
        raise Exception('Failed to fetch the current postgres version')

    # Grep the major postgres version
    str_arr = out.split(" ")
    major_ver = str_arr[2].split(".")[0]
    if major_ver is None:
        raise Exception('Failed to fetch the current postgres version')
    logger.info('Current major version: %s' % major_ver)
    return major_ver


def get_config_data():
    # Adding local import of vtsdb_utils as it is packaged in vstats RPM.
    # If we add this as global import it will cause upgrade failure in
    # precheck when upgrading from source machine where vstats
    # is not installed to a target where vstats is enabled as the
    # corresponding vstats modules would not be available during precheck stage.
    sys.path.append(os.sep.join([get_env_var('VMWARE_CIS_HOME'), 'vmware-vtsdb', 'scripts']))
    from vtsdb_utils import get_vtsdb_config_data
    # Dictionary to hold config params
    cfg = get_vtsdb_config_data()

    # postgres pg_upgrade entry
    cfg['PG_UPGRADE'] = os.sep.join([cfg['PG_BIN'], 'pg_upgrade'])

    # Postgres DB admin
    cfg['PG_DB_ADMIN'] = 'postgres'

    # vtsdb upgrade log
    cfg['VTSDB_UPGRADE_LOG'] = os.sep.join([cfg['VTSDB_LOG'], 'upgrade'])
    pg_ver = GetCurrentPGVersion()
    if pg_ver.isnumeric() and int(pg_ver) >= 10:
        cfg['PG_VERSION'] = pg_ver

    return cfg


def PGUpgradeInplaceRequired(cfg):
    """
    Method to check if major upgrade is required by
    comparing the version corresponding to the data files against the
    current PG_VERSION
    """
    logger.info("Checking if in place pg_upgrade is required")
    upg_src_ver = GetPgVersion(cfg['VTSDB_DATA'])

    upg_dest_ver = cfg['PG_VERSION']
    if upg_src_ver == upg_dest_ver:
        logger.info(
            'Major upgrade not needed, upgrade source version %s same as this version %s' % (upg_src_ver, upg_dest_ver))
        return False
    else:
        return True


# Run post-upgrade checks.
# This is aimed at preforming post-upgrade checks on an upgraded instance.
def VerifyUpgradeInplace(cfg):
    # Check that service is able to start properly. This code path
    # should as well stop the service before leaving.
    service_start_cmd = ['/bin/service-control', '--start', COMPONENT_NAME]
    service_stop_cmd = ['/bin/service-control', '--stop', COMPONENT_NAME]

    logger.info("Start the vtsdb service..")
    cmd = ['/bin/bash', '-c', '%s' % ' '.join(service_start_cmd)]
    logger.info('Running %s' % ' '.join(cmd))
    (ret, out, err) = run_command(cmd)
    if ret != 0:
        logger.error('vtsdb service start failed error code %d with error %s out %s' % (ret, err, out))
        return ret

    logger.info("Stop the vtsdb service..")
    cmd = ['/bin/bash', '-c', '%s' % ' '.join(service_stop_cmd)]
    logger.info('Running %s' % ' '.join(cmd))
    (ret, out, err) = run_command(cmd)
    if ret != 0:
        logger.error('vtsdb service stop failed error code %d with error %s out %s' % (ret, err, out))
        return ret

    return 0


def UpgradeInPlacePost(cfg):
    """
    Perform a couple of post pg_upgrade steps
    1. Create directory to store temporary stat
    2. Create socket directory for VTSDB instance (/var/run/VTSDB/.s.PGSQL.5433)
    3. Do the SCA registration
    :param cfg:
    :return:
    """
    # Adding local import of vtsdb_utils as it is packaged in vstats RPM.
    # If we add this as global import it will cause upgrade failure in precheck
    # when upgrading from source machine where vstats is not installed
    # to a target where vstats is enabled, as the corresponding vstats modules
    # would not be available during precheck stage.
    sys.path.append(os.sep.join([get_env_var('VMWARE_CIS_HOME'), 'vmware-vtsdb', 'scripts']))
    from vtsdb_utils import SetupPgConf, \
        CreateTempStatDir, CreateSocketDir, RegisterWithSCA
    logger.info("Running post actions of in-place pg_upgrade")

    ret = CreateTempStatDir(cfg)
    if ret == 0:
        logger.info('Created temp statistics dirs for VTSDB')
        ret = CreateSocketDir(cfg)
    if ret == 0:
        logger.info('Created socket directory for VTSDB')
        SetupPgConf(cfg)
    if ret == 0:
        logger.info('Created conf files for VTSDB')
        ret = RegisterWithSCA(cfg)
    if ret == 0:
        logger.info('Registered With SCA')

    return ret


def PGUpgradeInstanceInplace(cfg):
    """
    Run major upgrade
    :param cfg:
    :return:
    """
    # Adding local import of vpostgres_cis.upgrade as it is available from REL_12_STABLE.
    # If we add this as global import it will cause upgrade failure in precheck
    # when upgrading from source machine where vpostgres PG_VERSION
    # is less than 12.
    from vpostgres_cis.upgrade import InitDBForUpgrade, PgUpgradeInplace
    # Init db for new version
    logger.info("Initialize new data directory using initdb")
    ret = InitDBForUpgrade(data_dir=cfg['VTSDB_DATA'],
                           pg_base_path=VTSDB_HOME,
                           os_admin=cfg['VTSDB_OS_USER'],
                           initdb_opts=cfg['PG_INITDB_OPTS'])
    # Invoke pg_upgrade
    if ret == 0:
        conf_files = ['health_status_worker.conf']
        logger.info("Perform pg_upgrade in place")
        ret = PgUpgradeInplace(data_dir=cfg['VTSDB_DATA'],
                               xlog_dir=cfg['VTSDB_WAL_LOG'],
                               pg_base_path=VTSDB_HOME,
                               upgrade_log_dir=cfg['VTSDB_UPGRADE_LOG'],
                               os_admin=cfg['VTSDB_OS_USER'],
                               os_group=cfg['VTSDB_OS_GROUP'],
                               db_admin=cfg['PG_DB_ADMIN'],
                               conf_files=conf_files)
    if ret == 0:
        ret = UpgradeInPlacePost(cfg)
    # Verify upgrade
    if ret == 0:
        ret = VerifyUpgradeInplace(cfg)
    if ret == 0:
        logger.info('Major in-place upgrade to version %s was successful' % cfg['PG_VERSION'])

    return ret


