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

import os
import sys
import logging
import psycopg2
from lxml import etree
import psutil

logger = logging.getLogger(__name__)
from l10n import msgMetadata as _T, localizedString as _

rootDir = os.path.abspath(os.path.dirname(__file__))
sys.path.append(os.path.join(rootDir, '../../../', 'libs'))
sys.path.append(os.path.join(rootDir, '../../../', 'libs', 'feature-state'))
sys.path.append(os.path.join(rootDir, '../../../', 'libs', 'sdk'))

"""
Since precheck.py is consumed by both b2b and day0 patching which have
different importing cadence, need to add try/catch here for importing.
"""
try:
    from . import utils
except Exception:
    import utils

try:
    from vmware.update.specs import Mismatch
except Exception:
    from patch_specs import Mismatch

from . import autoupgrade_precheck

ID_ERR = 'com.vmware.precheck.error'
ID_ERR_DESC = 'com.vmware.precheck.error.description'
ID_ERR_AUTOUPGRADE = 'com.vmware.autoupgrade_precheck.error'


ID_ERR_DESC_AUTOUPGRADE = _T(
    'com.vmware.autoupgrade_precheck.error.description',
    'The vCenter Server upgrade is blocked because this Supervisor Cluster '
    'is running on an incompatible Kubernetes version : \n '
    '{ %(0)s:%(1)s }')

ID_RESOL_ERR_AUTOUPGRADE = _T(
    'com.vmware.autoupgrade_precheck.error.resolution',
    'Before you upgrade the vCenter Server, upgrade the Kubernetes control '
    'plane on the Supervisor Cluster to a supported version. '
    'The supported Kubernetes versions are : [%(0)s]')


ID_WARN_LICENSE = 'com.vmware.autoupgrade_license.warning'
ID_WARN_DESC_LICENSE = _T(
    'com.vmware.autoupgrade_license.warning.description',
    'The Tanzu edition license or the evaluation period '
    'of this Supervisor Cluster has expired: \n'
    '{ %(0)s:%(1)s }')
ID_RESOL_WARN_LICENSE_INCOMPATIBLE = _T(
    'com.vmware.autoupgrade_license_incompatible.warning.resolution',
    'Before proceeding with the upgrade, assign a valid Tanzu edition license '
    'to the Supervisor Cluster. You can upgrade the vCenter Server without '
    'assigning a valid license to the Supervisor Cluster. '
    'However, if any errors occur after the upgrade, you cannot '
    'recover the cluster or assign a valid Tanzu edition license to it.'
    )
ID_RESOL_WARN_LICENSE_COMPATIBLE = _T(
    'com.vmware.autoupgrade_license_compatible.warning.resolution',
    'Before proceeding with the upgrade, assign a valid Tanzu edition '
    'license to the Supervisor Cluster. You can upgrade the vCenter Server '
    'without assigning a valid license to the Supervisor Cluster. '
    'However, you must upgrade the Supervisor Cluster manually '
    'after the vCenter Server upgrade.')


ID_WARN_AUTOUPGRADE = 'com.vmware.autoupgrade_precheck.warning'
ID_WARN_DESC_AUTOUPGRADE = _T(
    'com.vmware.autoupgrade_precheck.warning.description',
    'This Supervisor Cluster is running an unsupported Kubernetes version '
    'and becomes incompatible after the vCenter Server upgrade : \n '
    '{ %(0)s:%(1)s }')

ID_RESOL_WARN_AUTOUPGRADE = _T(
    'com.vmware.autoupgrade_precheck.warning.resolution',
    'Upgrade Kubernetes on the Supervisor Cluster manually '
    'to one of these versions: [%(0)s]. '
    'If you do not upgrade Kubernetes manually, it is automatically '
    'upgraded to the next compatible version. '
    'Before you upgrade the Kubernetes control plane, '
    'ensure that the Supervisor Cluster is in the READY state.')

ID_RESOL = 'com.vmware.precheck.error.resolution'
ID_WARN = 'com.vmware.precheck.warning'
ID_WARN_DESC = 'com.vmware.precheck.warning.description'
ID_RESOLUTION_WARN = 'com.vmware.precheck.warning.resolution'

DESC_ERR = "There are clusters doing upgrade(not ready yet): "
TEXT_ERR = "There are clusters doing upgrade now(not ready yet), VC patch \
stops."
RESOL = "Retry VC patch after all on-going upgrade for wcp enabled clusters \
completed."

DESC_LIMIT_WARN_AUTOUPGRADE = _T(
    'com.vmware.autoupgrade_precheck.desc_limit.warning',
    'There are %(0)s more clusters with auto-upgrade errors and warnings.')
RESOLUTION_LIMIT_WARN_AUTOUPGRADE = _T(
    'com.vmware.autoupgrade_precheck.desc_limit.resolution',
    'Check the vCenter Server logs to view all '
    'auto-upgrade errors and warnings.')


# VKAL-4295, we are querying VCDB directly to gather the information about the
# WCP clusters instead of using WCP APIs in case the APIs are unavailable.
# The database view cluster_upgrade_state is being created and refreshed in
# wcpsvc.
WCP_PGPASSFILE = "/etc/vmware/wcp/.pgpass"
# This is a db view created as VKAL-4295
DB_VIEW_CLUSTER_UPGRADE_STATE = "cluster_upgrade_state"
# This is the fallback table name. cluster_upgrade_state view was created
# from this table. To support backward-compatibility to cover the case where
# the cluster_upgrade_state view was not present yet.
DB_TABLE_CLUSTER_DB_CONFIG = "cluster_db_configs"
DB_QUERY_GET_CLUSTER_STATE = "select cluster, upgrade_state, cluster_version" \
                             ", desired_version, upgrade_task from {};"
VCHA_CONFIG_FILEPATH = "/etc/vmware-vcha/vcha.cfg"


class WcpClusterUpgradeState:
    """
    A data structure that holds upgrade state of WCP supervisor cluster that
    is retrieved from WCP database.
    """

    def __init__(self):
        self.cluster = ''
        self.upgrade_state = ''
        self.cluster_version = ''
        self.desired_version = ''
        self.upgrade_task = ''

    def __repr__(self):
        return self.__str__()

    def __str__(self):
        return str(self.__dict__)


def execute_vcdb_query(query):
    """Execute query against VCDB and return List of tuple
    :return: List of tuple
    :rtype: list[tuple]
    """
    conn = cur = None
    os.environ['PGPASSFILE'] = WCP_PGPASSFILE
    try:
        conn = psycopg2.connect("host=localhost dbname=VCDB user=wcpuser")
        cur = conn.cursor()
        cur.execute(query)
        results = cur.fetchall()  # list[tuple(cluster, upgrade_state,
        # cluster_version, desired_version, upgrade_task)]
        return results
    except Exception as e:
        logger.error("Failed to query VCDB for cluster upgrade state. Err {}"
                     .format(str(e)))
        return []
    finally:
        if cur:
            cur.close()
        if conn:
            conn.close()


def query_cluster_upgrade_state():
    """Query the cluster upgrade state from VCDB directly, first we try to query
    if from DB_VIEW_CLUSTER_UPGRADE_STATE, if it fails (maybe the view is not
    yet created), then we fall back to query DB_TABLE_CLUSTER_DB_CONFIG.

    :return: List of WcpClusterUpgradeState
    :rtype: list[WcpClusterUpgradeState]
    """
    # We first check from cluster_upgrade_state view, if this view does not
    # exist, then fallback to use cluster_db_configs table.
    table_names = [DB_VIEW_CLUSTER_UPGRADE_STATE, DB_TABLE_CLUSTER_DB_CONFIG]
    for table_name in table_names:
        query = DB_QUERY_GET_CLUSTER_STATE.format(table_name)
        results = execute_vcdb_query(query)  # type: list[tuple]
        ret = []
        for t in results:
            obj = WcpClusterUpgradeState()
            obj.cluster = t[0]
            obj.upgrade_state = t[1]
            obj.cluster_version = t[2]
            obj.desired_version = t[3]
            obj.upgrade_task = t[4]
            ret.append(obj)
        if ret:
            return ret
    logger.debug("No WCP-enabled cluster found, proceed with patching.")
    return []


def get_ip_addresses():
    """Return all ip addresses from all nic for this machine.

    :return: All IPs from all nics
    :rtype: list[str]
    """
    nics = psutil.net_if_addrs()
    ret = []
    for nic in nics.values():
        ret.extend([x.address for x in nic])
    return ret


def is_vcha_witness_node():
    """Return true if this VC that's running this script is a VCHA witness node.
    This will check VCHA_CONFIG_FILEPATH for witness's ip address and see if it
    matches any of the IPs this node has.

    :return: true if this VC is a VCHA witness node
    :rtype: bool
    """
    try:
        with open(VCHA_CONFIG_FILEPATH) as f:
            content = f.read()
    except:
        # no such file or can't read the file, treat it as False
        logger.debug("File {} does not exist, consider it not a VCHA witness "
                     "node".format(VCHA_CONFIG_FILEPATH))
        return False
    root = etree.fromstring(content)
    config = etree.ElementTree(root)
    node = config.find(".//witnessIP")
    if node is None:
        logger.warning("There is no witnessIP found in {}".format(
            VCHA_CONFIG_FILEPATH))
        return False
    witness_ip = node.text.strip()  # type: str
    current_ips = get_ip_addresses()
    if witness_ip in current_ips:
        return True
    logger.debug("VCHA witness ip {0} does not match any current IP, this is"
                 "not a VCHA witness node.".format(witness_ip))
    return False


def generateMismatchforAutouprade(auto_upgrade_clusters, supported_versions):
    """
    Function to generate mismatch which could be either warning or error,
    which would be displayed to the customer
    The function loops over the list of clusters returned after
    autoupgrade precheck and returns appropriate warning or error
    messages complying to the cluster
    """
    mismatches = []
    for item in auto_upgrade_clusters:
        if item.compatibility:
            logger.info("Cluster lies in n-3 window,"
                        "returning warning message")
            cluster_details = [item.cluster, item.cluster_version]
            if not item.license_status:
                # If cluster has invalid license, VC upgrade is allowed to go
                # through with a warning
                logger.info("Cluster has invalid license,"
                            "returning warning message")
                mismatch = Mismatch(
                    text=_(ID_WARN_DESC_LICENSE, cluster_details),
                    problemId=ID_WARN_LICENSE,
                    description=_(ID_WARN_DESC_LICENSE, cluster_details),
                    resolution=_(ID_RESOL_WARN_LICENSE_COMPATIBLE),
                    severity=Mismatch.WARNING)
                mismatches.append(mismatch)
            else:
                logger.info("cluster has a valid license")
                supported_version_list = ', '.join(
                    str(e) for e in supported_versions)
                mismatch = Mismatch(
                    text=_(ID_WARN_DESC_AUTOUPGRADE, cluster_details),
                    problemId=ID_WARN_AUTOUPGRADE,
                    description=_(ID_WARN_DESC_AUTOUPGRADE, cluster_details),
                    resolution=_(ID_RESOL_WARN_AUTOUPGRADE,
                                 supported_version_list),
                    severity=Mismatch.WARNING)
                mismatches.append(mismatch)
        else:
            logger.info("Cluster lies in n-4 window,"
                        "returning error message")
            cluster_details = [item.cluster, item.cluster_version]
            if not item.license_status:
                # If cluster has invalid license, VC upgrade is allowed to go
                # through with a warning
                logger.info("Cluster has invalid license,"
                            "returning warning message")
                mismatch = Mismatch(
                    text=_(ID_WARN_DESC_LICENSE, cluster_details),
                    problemId=ID_WARN_LICENSE,
                    description=_(ID_WARN_DESC_LICENSE, cluster_details),
                    resolution=_(ID_RESOL_WARN_LICENSE_INCOMPATIBLE),
                    severity=Mismatch.WARNING)
                mismatches.append(mismatch)
            else:
                supported_version_list = ', '.join(
                    str(e) for e in supported_versions)
                mismatch = Mismatch(
                    text=_(ID_ERR_DESC_AUTOUPGRADE, cluster_details),
                    problemId=ID_ERR_AUTOUPGRADE,
                    description=_(ID_ERR_DESC_AUTOUPGRADE, cluster_details),
                    resolution=_(ID_RESOL_ERR_AUTOUPGRADE,
                                 supported_version_list),
                    severity=Mismatch.ERROR)
                mismatches.append(mismatch)
    logger.info("mismatches from auto-upgrade prechecks:{}".format(mismatches))
    if len(mismatches) > 10:
        pending_clusters = len(mismatches) - 10
        # Limiting the number of mismatches returned by auto-upgrade to 10
        mismatches = mismatches[0:10]
        mismatch = Mismatch(
            text=_(DESC_LIMIT_WARN_AUTOUPGRADE, str(pending_clusters)),
            problemId=ID_WARN_AUTOUPGRADE,
            description=_(DESC_LIMIT_WARN_AUTOUPGRADE, str(pending_clusters)),
            resolution=_(RESOLUTION_LIMIT_WARN_AUTOUPGRADE),
            severity=Mismatch.WARNING)
        mismatches.append(mismatch)
    return mismatches


def doPrecheck(skipPrechecksIds):
    """
    Make precheck as a module consumed by both b2b and day0 VC patching.
    Here are the prechecks conducted:
    1. Check if WCP is installed on VC. If not, skip prechecks.
    2. Check if there are WCP enabled clusters. If no, skip prechecks.
    3. Check if there are clusters now doing upgrade when VC patching.
    If yes, raise error in Mismatch.
    4. Check if there are clusters running k8s version not supported after
    VC patching. If yes, conduct auto upgrade prechecks on the clusters.
    Moreover, option for skipping prechecks are provided by users via
    PrechecksIds.
    If user specified wcp.cluster.upgrading, then precheck #3 is skipped.
    If user specified wcp.cluster.k8s_check, then precheck #4 is skipped.

    :param skipPrechecksIds
    :type skipPrechecksIds: list([str])
    :return: list of Mismatch found in prechecks
    :rtype: list([Mismatch])
    """
    mismatch = None
    mismatches = []
    skipUpgrade = False
    skipK8sCheck = False
    skipUpgrade = "wcp.cluster.upgrading" in skipPrechecksIds
    skipK8sCheck = "wcp.cluster.k8s_check" in skipPrechecksIds

    # skip all prechecks if this VC is a VCHA witness node since there is no
    # local VC state, VCDB, and etc.
    if is_vcha_witness_node():
        logger.debug("The current node is a VCHA witness node, skip precheck.")
        return mismatches

    # Check if WCP is installed on VC. If not, skip prechecks.
    if not os.path.exists(utils.__VERSION_FILE__):
        logger.info("WCP is not enabled on this VC")
        return mismatches

    # Check if any new version is being introduced
    if not utils.patch_payload_has_new_k8s_versions():
        logger.info('Incoming patch does not introduce new k8s versions, '
                    'at-risks cluster check skipped.')
        return mismatches

    # Query VCDB to get all wcp enabled clusters. We are not querying this using
    # wcp API because during VC patching, the wcpsvc might be down and we can't
    # query it.
    wcp_enabled_clusters = query_cluster_upgrade_state()
    logger.debug("WCP enabled clusters: {}".format(wcp_enabled_clusters))

    if not wcp_enabled_clusters:
        logger.info("There is no WCP enabled clusters deployed on this VC")
        return mismatches

    if not skipUpgrade:
        # There are three states for wcp enabled cluster:
        # READY: Clusters finished upgrade or ready to do upgrade.
        # ERROR: Clusters FAILED in upgrade.
        # PENDING: Clusters are now doing upgrade.
        # Check if there are wcp enabled clusters now doing upgrade by
        # checking if Clusters.State equals to PENDING.
        pending_clusters = [x for x in wcp_enabled_clusters
                            if x.upgrade_state == "PENDING"]
        # If not empty, check if there are clusters doing upgrade.
        if len(pending_clusters) > 0:
            desc_err = DESC_ERR + \
                       " ".join([s.cluster for s in pending_clusters])
            mismatch = Mismatch(
                text=_(_T(ID_ERR, desc_err)),
                problemId=ID_ERR,
                description=_(_T(ID_ERR_DESC, desc_err)),
                resolution=_(_T(ID_RESOL, RESOL)),
                severity=Mismatch.ERROR)

            mismatches.append(mismatch)
            return mismatches
    else:
        logger.info("Pre-check for on-going upgrade cluster skipped")

    if not skipK8sCheck:
        logger.info("Checking for at risk clusters")
        # check if the current version is present in the incoming patch
        at_risk_clusters, supported_versions = utils.get_at_risk_clusters(wcp_enabled_clusters)
        if len(at_risk_clusters) > 0:
            logger.info("at risk clusters present, conducting prechecks "
                        "for autoupgrade")
            try:
                auto_upgr_clusters = autoupgrade_precheck. \
                    doAutoUpgradePrecheck(at_risk_clusters)
                logger.info("Clusters received after autoupgrade prechecks:%s",
                            auto_upgr_clusters)
                mismatches = generateMismatchforAutouprade\
                    (auto_upgr_clusters, list(supported_versions))
                return mismatches
            except Exception:
                msg = "Error while conducting prechecks"
                logger.info(msg)
                raise Exception(msg)
    else:
        logger.info("Pre-check for Kubernetes version compatibility skipped.")
    logger.info("WCP precheck PASSED")
    return mismatches
