# Copyright (c) 2019-2021 VMware, Inc.  All rights reserved.
# -- VMware Confidential
# coding: utf-8

import os
import sys
import shutil
import yaml
import filecmp
import logging

# VMWARE_PYTHON_PATH contains cis.tools and pyVim module, however not every
# caller of this script adds it to python module path. We add the path
# automatically here to simplify callers of this script.
for entry in os.environ['VMWARE_PYTHON_PATH'].split(":"):
    if entry not in sys.path:
        sys.path.append(entry)

logger = logging.getLogger(__name__)

CONFIG_DIR = "/etc/vmware/wcp/"
VAR_DIR = "/var/vmware/wcp"
# Temp marker file generated by b2b pre-patch hook for marking that b2b
# patching framework is in progress and suppressing the triggering of
# of patchingOrchestrator in post-rpm and wcp-prestart script.
WCP_B2B_PATCHING_IN_PROGRESS = "/storage/wcp/b2b_patching_in_progress"

# Temp marker file generated by patchingOrchestrator triggered in rpm
# post-install script and wcp service pre-start script. Since patching
# can not take place during wcp service normal starts/restarts, it is needed
# to generate a temp marker file to notify wcp pre-start script whether it is
# in patching process or not.
# Both B2B and DAY0 marker file are generated in directory owned/dedicated by
# wcpsvc.
WCP_DAY0_PATCHING_IN_PROGRESS = "/storage/wcp/wcp_day0_patching_in_progress"

# The file wcp_rdu_in_progress_src will be created on the source VC
# during the expand hook of RDU upgrade. It will be copied over to the
# target VC as file wcp_rdu_in_progress to indicate to the patchingOrchestrator
# that we are in the process of performing RDU upgrade.
WCP_RDU_IN_PROGRESS = "/storage/wcp/wcp_rdu_in_progress"
WCP_RDU_IN_PROGRESS_SRC = "/storage/wcp/wcp_rdu_in_progress_src"

# The file will be installed as part of the rpm. During upgrade, we need
# to read its original content before the rpm overwrites the original file,
# e.g. @Discovery hook.
__VERSION_FILE__ = os.path.join(CONFIG_DIR, "wcp_version.txt")
__VERSION_FILE_NAME__ = os.path.basename(__VERSION_FILE__)
PRESERVED_VERSION_FILE_NAME = "wcp_version_preserved.txt"
# vpxd-extension.xml for WCP defining WCP task ids that needs to be registered
# after VC patching.
__VPXD_EXTENSION_FILE__ = os.path.join(CONFIG_DIR, "vpxd-extension.xml")
__VPXD_EXTENSION_FILE_NAME__ = os.path.basename(__VPXD_EXTENSION_FILE__)

MY_PAYLOAD_DIR = os.path.dirname(__file__)
PATCHES_DIR = os.path.join(MY_PAYLOAD_DIR, "patches")
VERSIONS_FILEPATH = os.path.join(MY_PAYLOAD_DIR, "patches.yaml")
PATCH_PAYLOAD_KUBERNETES_VERSIONS_FILEPATH = os.path.join(MY_PAYLOAD_DIR,
                                                          "wcp_versions.yaml")
VC_KUBERNETES_VERSIONS_FILEPATH = os.path.join(CONFIG_DIR, "wcp_versions.yaml")

SAVED_WCPSVC_YAML = os.path.join(CONFIG_DIR, "wcpsvc.yaml.save")
WCPSVC_YAML = "wcpsvc.yaml"


def getCurrentVersion(stage_dir, file_name=__VERSION_FILE_NAME__):
    """Return the version from the staged version file inside
    stage Directory if any. Return stripped string or None if not
    found.

    :param stage_dir: Directory to preserve wcp version file.
    :type stage_dir: str
    :param file_name: The file name of wcp version file.
    :type file_name: str
    :rtype: str
    """
    file_path = os.path.join(stage_dir, file_name)
    if os.path.exists(file_path):
        with open(file_path) as fp:
            version = fp.read().strip()
    else:
        version = None
    return version


def getStagedWcpsvcYamlPath(stage_dir):
    """
    Return the staged version of wcpsvc.yaml
    """
    wcpsvc_yaml_filepath = os.path.join(MY_PAYLOAD_DIR, WCPSVC_YAML)
    return os.path.join(stage_dir, wcpsvc_yaml_filepath)


def cmpDiff(old_file, new_file):
    """
    Compare the difference between two file

    :param old_file: File path of the old file.
    :type old_file: str
    :param new_file: File path of the new file.
    :type new_file: str
    :rtype: bool
    """
    if not os.path.exists(old_file):
        return True
    return not filecmp.cmp(old_file, new_file)


def preserveFile(stage_dir, file_path=__VERSION_FILE__):
    """ Persists certain file in the stage directory.

    :param stage_dir: Directory to preserve wcp version file.
    :type stage_dir: str
    :param file_path: The filepath wcp version file to be preserved.
    :type file_path: str
    """
    if os.path.exists(file_path):
        shutil.copy(file_path, stage_dir)


def getApplicableSortedPatches(file_path=VERSIONS_FILEPATH):
    """
    Read from version2patch.yaml and return list of all patches.

    :param file_path: directory where version2yaml file locates
    :return: list of patches
    :rtype: list([str])
    """
    with open(file_path) as fp:
        patches = yaml.safe_load(fp)

    return patches


def extractVersion(version):
    """Return only major.minor version of input k8s version.
    :param version: k8s version, in the form of major.minor.patch.
    :return: k8s version with only major.minor format.
    :rtype: str
    """
    if version[0].lower() == 'v':
        version = version[1:]
    return version.split(".")[0] + "." + version.split(".")[1]


def extract_k8s_versions(filepath):
    """Extracts and returns list of Kubernetes major.minor versions
    (stripping off patch version, since patch versions do not impact
    upgrade compatibility) from given WCP version manifest.

    :type filepath: string
    :return: list of Kubernetes version strings
    :rtype: list
    """
    logger.debug('Extracting Kubernetes versions from %s', filepath)
    try:
        # Expected yaml doc structure:
        # wcp_versions:
        # - supported_k8s_versions:
        #   - version: v1.16.2
        #     ...
        #   - version: v1.15.4
        #     ...
        # ...
        with open(filepath) as f:
            doc = yaml.safe_load(f)
        wcp_versions = doc['wcp_versions'][0]
        all_versions = wcp_versions['supported_k8s_versions']
        supported_k8s = set([extractVersion(v['version'])
                             for v in all_versions])
        logger.info('Found Kubernetes versions %s in %s', supported_k8s,
                    filepath)
        return supported_k8s
    except Exception:
        msg = 'Failed to load Kubernetes version data from %s' % filepath
        logger.exception(msg)
        raise Exception(msg)


def patch_payload_has_new_k8s_versions():
    """Returns true if incoming patch payload has newer Kubernetes
    major.minor versions compared to existing Kubernetes versions supported
    in VC's control plane VM ova image.

    :return: true if incoming patch has new Kubernetes versions
    :rtype: bool
    """
    incoming_vers_file = PATCH_PAYLOAD_KUBERNETES_VERSIONS_FILEPATH
    incoming_k8s_versions = extract_k8s_versions(incoming_vers_file)
    existing_vers_file = VC_KUBERNETES_VERSIONS_FILEPATH
    existing_k8s_versions = extract_k8s_versions(existing_vers_file)
    return incoming_k8s_versions != existing_k8s_versions


def upgrade_version_compatibility_check():
    """
    This function checks if the current version is present in
    "supported upgrade from" list in wcp_versions.yaml
    It returns a list of clusters that are autoupgrade
    compatible; if the current cluster version exists in the
    list, it is "n-3" and can be autoupgraded
    """
    logger.info("getting autoupgrade compatible versions")
    supported_upgrade_versions = []
    with open(PATCH_PAYLOAD_KUBERNETES_VERSIONS_FILEPATH) as fp:
        wcp_versions = yaml.safe_load(fp)['wcp_versions'][0]

    all_versions = wcp_versions['supported_k8s_versions']
    for versions in all_versions:
      supported_k8s_versions = versions['supported_upgrade_from']
      for element in supported_k8s_versions:
        if element not in supported_upgrade_versions:
          supported_upgrade_versions.append(element)

    logger.info("all the upgrade commpatible versions returned %s:",
                supported_upgrade_versions)
    return supported_upgrade_versions


def get_at_risk_clusters(clusters):
    """Parse the wcp_versions.yaml file which contains the supported k8s
    version after wcpsvc upgrade. Check if there are clusters running "n-3"
    k8s version(at-risk clusters) which falls off the "n-2" k8s window after
    VC patch.
    The function returns the list of clusters that may be in n-3 or more window

    :type clusters: list[WcpClusterUpgradeState]
    :return: list of "at-risk" clusters.
    :rtype: list[WcpClusterUpgradeState]
    """
    logger.info("Getting at risk cluster")
    with open(PATCH_PAYLOAD_KUBERNETES_VERSIONS_FILEPATH) as fp:
        wcp_versions = yaml.safe_load(fp)['wcp_versions'][0]

    all_versions = wcp_versions['supported_k8s_versions']
    supported_k8s = set([extractVersion(v['version']) for v in all_versions])
    logger.info("All supported versions %s", supported_k8s)
    at_risk_clusters = [x for x in clusters if
                        extractVersion(x.cluster_version) not in supported_k8s]

    logger.info("at risk clusters returned %s:", at_risk_clusters)
    return at_risk_clusters, supported_k8s