# Copyright (c) 2015-2016 Tigera, Inc. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import hashlib
import json
import logging
import subprocess
import time
import types
from multiprocessing.dummy import Pool as ThreadPool
from pprint import pformat
from unittest import TestCase

import yaml
from deepdiff import DeepDiff

from tests.st.utils.utils import (get_ip, ETCD_SCHEME, ETCD_CA, ETCD_CERT,
                                  ETCD_KEY, debug_failures, ETCD_HOSTNAME_SSL,
                                  wipe_etcd, clear_on_failures)

# The number of test batches used in CI.
NUM_BATCHES = 6

HOST_IPV6 = get_ip(v6=True)
HOST_IPV4 = get_ip()

logging.basicConfig(level=logging.DEBUG, format="%(message)s")
logger = logging.getLogger(__name__)

# Disable spammy logging from the sh module
sh_logger = logging.getLogger("sh")
sh_logger.setLevel(level=logging.CRITICAL)

first_log_time = None


def calculate_batch(cname, mname):
    combined = cname + "." + mname
    m = hashlib.sha224()
    m.update(combined)
    batch = ord(m.digest()[0]) % NUM_BATCHES
    print "Assigned %s to batch %s" % (combined, batch)
    return batch


class AutoBatcher(type):
    """
    AutoBatcher is a metaclass that makes sure every test_ method has a batchnumber.

    Batch numbers are assigned deterministically using a hash of class name and
    method name.
    """

    def __init__(cls, name, bases, dct):
        test_methods_with_no_batch = {}
        has_batch = False
        for k, v in dct.iteritems():
            if k == "batchnumber":
                has_batch = True
            elif k.startswith("test_"):
                if not isinstance(v, types.FunctionType):
                    continue
                if hasattr(v, "batchnumber"):
                    continue
                test_methods_with_no_batch[k] = v
        if not has_batch:
            for k, v in test_methods_with_no_batch.iteritems():
                v.batchnumber = calculate_batch(name, k)
            dct["batchnumber"] = calculate_batch(name, "__class__")
        super(AutoBatcher, cls).__init__(name, bases, dct)


class TestBase(TestCase):
    __metaclass__ = AutoBatcher

    """
    Base class for test-wide methods.
    """
    @classmethod
    def setUpClass(cls):
        wipe_etcd(HOST_IPV4)

    def setUp(self, clear_etcd=True):
        """
        Clean up before every test.
        """
        self.ip = HOST_IPV4

        if clear_etcd:
            wipe_etcd(self.ip)

        # Log a newline to ensure that the first log appears on its own line.
        logger.info("")

        clear_on_failures()

    @staticmethod
    def _conn_checker(args):
        source, dest, test_type, result, retries = args
        if test_type == 'icmp':
            if result:
                return source.check_can_ping(dest, retries)
            else:
                return source.check_cant_ping(dest, retries)
        elif test_type == 'tcp':
            if result:
                return source.check_can_tcp(dest, retries)
            else:
                return source.check_cant_tcp(dest, retries)
        elif test_type == 'udp':
            if result:
                return source.check_can_udp(dest, retries)
            else:
                return source.check_cant_udp(dest, retries)
        else:
            logger.error("Unrecognised connectivity check test_type")

    @debug_failures
    def assert_connectivity(self, pass_list, fail_list=None, retries=0,
                            type_list=None):
        """
        Assert partial connectivity graphs between workloads.

        :param pass_list: Every workload in this list should be able to ping
        every other workload in this list.
        :param fail_list: Every workload in pass_list should *not* be able to
        ping each workload in this list. Interconnectivity is not checked
        *within* the fail_list.
        :param retries: The number of retries.
        :param type_list: list of types to test.  If not specified, defaults to
        icmp only.
        """
        if type_list is None:
            type_list = ['icmp', 'tcp', 'udp']
        if fail_list is None:
            fail_list = []

        # Wait (up to 20 retries) for each of the workloads to be able to ping
        # _itself_.  That will ensure that each workload is properly up and
        # running.
        self.wait_assert_self_connectivity(pass_list + fail_list, type_list)

        # Now check that connectivity _between_ the workloads is as expected.
        conn_check_list = []
        for source in pass_list:
            for dest in pass_list:
                if source is dest:
                    continue
                if 'icmp' in type_list:
                    conn_check_list.append((source, dest.ip, 'icmp', True, retries))
                if 'tcp' in type_list:
                    conn_check_list.append((source, dest.ip, 'tcp', True, retries))
                if 'udp' in type_list:
                    conn_check_list.append((source, dest.ip, 'udp', True, retries))
            for dest in fail_list:
                if 'icmp' in type_list:
                    conn_check_list.append((source, dest.ip, 'icmp', False, retries))
                if 'tcp' in type_list:
                    conn_check_list.append((source, dest.ip, 'tcp', False, retries))
                if 'udp' in type_list:
                    conn_check_list.append((source, dest.ip, 'udp', False, retries))

        results, diagstring = self.probe_connectivity(conn_check_list)
        assert False not in results, ("Connectivity check error!\r\n"
                                      "Results:\r\n %s\r\n" % diagstring)

    def wait_assert_self_connectivity(self, workloads, type_list):
        # First check that each of the workloads that we're looking at can ping
        # _itself_, allowing 20 retries for that.
        self_conn_list = []
        for source in workloads:
            if 'icmp' in type_list:
                self_conn_list.append((source, source.ip, 'icmp', True, 20))
            if 'tcp' in type_list:
                self_conn_list.append((source, source.ip, 'tcp', True, 20))
            if 'udp' in type_list:
                self_conn_list.append((source, source.ip, 'udp', True, 20))

        results, diagstring = self.probe_connectivity(self_conn_list)
        assert False not in results, ("Self-connectivity check error!\r\n"
                                      "Results:\r\n %s\r\n" % diagstring)

    def probe_connectivity(self, conn_check_list):
        # Empirically, 18 threads works well on my machine!
        check_pool = ThreadPool(18)
        results = check_pool.map(self._conn_checker, conn_check_list)
        check_pool.close()
        check_pool.join()
        # _conn_checker should only return None if there is an error in calling it
        assert None not in results, ("_conn_checker error - returned None")
        diagstring = ""
        # Check that all tests passed
        if False in results:
            # We've failed, lets put together some diags.
            header = ["source.ip", "dest.ip", "type", "expected", "actual"]
            diagstring = "{: >18} {: >18} {: >7} {: >6} {: >6}\r\n".format(*header)
            for i in range(len(conn_check_list)):
                source, dest, test_type, exp_result, retries = conn_check_list[i]
                pass_fail = results[i]
                # Convert pass/fail into an actual result
                if not pass_fail:
                    actual_result = not exp_result
                else:
                    actual_result = exp_result
                diag = [source.ip, dest, test_type, exp_result, actual_result]
                diagline = "{: >18} {: >18} {: >7} {: >6} {: >6}\r\n".format(*diag)
                diagstring += diagline

        return results, diagstring

    @debug_failures
    def assert_ip_connectivity(self, workload_list, ip_pass_list,
                               ip_fail_list=None, type_list=None, retries=0):
        """
        Assert partial connectivity graphs between workloads and given ips.

        This function is used for checking connectivity for ips that are
        explicitly assigned to containers when added to calico networking.

        :param workload_list: List of workloads used to check connectivity.
        :param ip_pass_list: Every workload in workload_list should be able to
        ping every ip in this list.
        :param ip_fail_list: Every workload in workload_list should *not* be
        able to ping any ip in this list. Interconnectivity is not checked
        *within* the fail_list.
        :param type_list: list of types to test.  If not specified, defaults to
        icmp only.
        """
        if type_list is None:
            type_list = ['icmp']
        if ip_fail_list is None:
            ip_fail_list = []

        # Wait (up to 20 retries) for each of the workloads to be able to ping
        # _itself_.  That will ensure, at least, that each _source_ workload is
        # properly up and running, before we start trying to ping from it to
        # another IP.  (Unfortunately we can't do anything here to ensure that
        # the target IPs are actually available.)
        self.wait_assert_self_connectivity(workload_list, type_list)

        conn_check_list = []
        for workload in workload_list:
            for ip in ip_pass_list:
                if 'icmp' in type_list:
                    conn_check_list.append((workload, ip, 'icmp', True, retries))
                if 'tcp' in type_list:
                    conn_check_list.append((workload, ip, 'tcp', True, retries))
                if 'udp' in type_list:
                    conn_check_list.append((workload, ip, 'udp', True, retries))

            for ip in ip_fail_list:
                if 'icmp' in type_list:
                    conn_check_list.append((workload, ip, 'icmp', False, retries))
                if 'tcp' in type_list:
                    conn_check_list.append((workload, ip, 'tcp', False, retries))
                if 'udp' in type_list:
                    conn_check_list.append((workload, ip, 'udp', False, retries))

        results, diagstring = self.probe_connectivity(conn_check_list)
        assert False not in results, ("Connectivity check error!\r\n"
                                      "Results:\r\n %s\r\n" % diagstring)

    def check_data_in_datastore(self, host, data, resource, yaml_format=True):
        if yaml_format:
            out = host.calicoctl(
                "get %s --output=yaml" % resource)
            output = yaml.safe_load(out)
        else:
            out = host.calicoctl(
                "get %s --output=json" % resource)
            output = json.loads(out)
        self.assert_same(data, output)

    @staticmethod
    def assert_same(thing1, thing2):
        """
        Compares two things.  Debug logs the differences between them before
        asserting that they are the same.
        """
        assert cmp(thing1, thing2) == 0, \
            "Items are not the same.  Difference is:\n %s" % \
            pformat(DeepDiff(thing1, thing2), indent=2)

    @staticmethod
    def writeyaml(filename, data):
        """
        Converts a python dict to yaml and outputs to a file.
        :param filename: filename to write
        :param data: dictionary to write out as yaml
        """
        with open(filename, 'w') as f:
            text = yaml.dump(data, default_flow_style=False)
            logger.debug("Writing %s: \n%s" % (filename, text))
            f.write(text)

    @staticmethod
    def writejson(filename, data):
        """
        Converts a python dict to json and outputs to a file.
        :param filename: filename to write
        :param data: dictionary to write out as json
        """
        with open(filename, 'w') as f:
            text = json.dumps(data,
                              sort_keys=True,
                              indent=2,
                              separators=(',', ': '))
            logger.debug("Writing %s: \n%s" % (filename, text))
            f.write(text)

    @debug_failures
    def assert_false(self, b):
        """
        Assert false, wrapped to allow debugging of failures.
        """
        assert not b

    @debug_failures
    def assert_true(self, b):
        """
        Assert true, wrapped to allow debugging of failures.
        """
        assert b

    @staticmethod
    def log_banner(msg, *args, **kwargs):
        global first_log_time
        time_now = time.time()
        if first_log_time is None:
            first_log_time = time_now
        time_now -= first_log_time
        elapsed_hms = "%02d:%02d:%02d " % (time_now / 3600,
                                           (time_now % 3600) / 60,
                                           time_now % 60)

        level = kwargs.pop("level", logging.INFO)
        msg = elapsed_hms + str(msg) % args
        banner = "+" + ("-" * (len(msg) + 2)) + "+"
        logger.log(level, "\n" +
                   banner + "\n"
                            "| " + msg + " |\n" +
                   banner)

# Add a default batch number to all tests (used for running tests in parallel)
TestBase.batchnumber = 0
