#!/usr/bin/python3

# This file is part of Cockpit.
#
# Copyright (C) 2015 Red Hat, Inc.
#
# Cockpit is free software; you can redistribute it and/or modify it
# under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation; either version 2.1 of the License, or
# (at your option) any later version.
#
# Cockpit is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with Cockpit; If not, see <http://www.gnu.org/licenses/>.

import argparse
import os
import sys
import subprocess
import datetime
import glob
import time

from testlib_avocado.machine_core import testvm
from testlib_avocado.machine_core.constants import TEST_OS_DEFAULT

BASE_IMAGE = os.environ.get("TEST_OS", TEST_OS_DEFAULT)
CURRENT_DIR = os.path.dirname(os.path.realpath(__file__))
if CURRENT_DIR not in sys.path:
    sys.path.insert(1, CURRENT_DIR)

TEST_TIMEOUT = 600

# this is where we publish logs on the local machine
arg_attachments_dir = os.environ.get("TEST_ATTACHMENTS", None)
if not arg_attachments_dir:
    arg_attachments_dir = os.getcwd()
if not os.path.exists(arg_attachments_dir):
    os.makedirs(arg_attachments_dir)


def own_print(*args, **kwargs):
    # try to write moretimes to workaround BlockingIOError and if failed continue to next line
    retry_count = 3
    for counter in range(retry_count):
        try:
            print(*args, **kwargs)
            break
        except BlockingIOError:
            time.sleep(0.1)


class TAPformatter:
    counter = 1
    counter_failed = 0
    counter_skipped = 0
    total = 0

    def __init__(self, test_amount):
        self.total = test_amount

    def tap_header(self):
        own_print("\nTAP output -------------------------")
        own_print("1..%d" % self.total)

    def tap_summary(self, duration, restarted_tests=None):
        if restarted_tests:
            own_print("# RESTARTED %d tests: %s" % (len(restarted_tests), restarted_tests))
        if self.counter_failed == 0:
            own_print("# TESTS PASSED duration: %ds" % duration)
        else:
            own_print("# %d TESTS FAILED duration: %ds" % (self.counter_failed, duration))

    @staticmethod
    def test_name_string(test_name):
        test_tuple = os.path.basename(test_name).split(":", 1)
        return "{} ({})".format(test_tuple[0], test_tuple[1])

    def tap_item_header(self, test_name):

        own_print("# ----------------------------------------------------------------------")
        own_print("# %s" % self.test_name_string(test_name))
        own_print("#")

    def tap_item_summary(self, test_name, output, duration):
        status = None
        rest_line = ""
        for line in output.splitlines()[::-1]:
            if test_name in line:
                status, rest_line = line.strip().split(" ", 1)
                break
        for line in output.splitlines():
            own_print(line)
        own_print("# TEST LOG END -------")
        if status == "SKIP":
            own_print("ok %s %s # duration: %ds # SKIP %s\n" % (self.counter, self.test_name_string(test_name), duration, rest_line))
            self.counter_skipped += 1
        elif status == "PASS":
            own_print("ok %s %s # duration: %ds" % (self.counter, self.test_name_string(test_name), duration))
        else:
            own_print("not ok %s %s # duration: %ds # %s" % (self.counter, self.test_name_string(test_name), duration, rest_line))
            self.counter_failed += 1
        self.counter += 1

    def get_status(self):
        return not bool(self.counter_failed)


class AvocadoTestSuite:
    target_test_dir = "/tmp/avocado_tests"
    target_library_path = "/tmp/avocado_library"
    library_name = "testlib_avocado"
    base_cmd = "python3 -m avocado"
    tap = None
    env = {}
    counter = 0
    timeout = 600
    failed_counter = 0
    machine_cockpit = None
    machine_avocado = None
    ipaddr_avocado = "10.111.113.200"
    ipaddr_cockpit = "10.111.113.1"

    def __init__(self, network, verbose, machine_cockpit=None):
        self.env = {"PYTHONPATH": self.target_library_path}
        self.machine_cockpit = machine_cockpit
        self.network = network
        self.verbose = verbose
        # store last executed avocado command to variable
        self.avocado_command = None
        self.machine_avocado = testvm.VirtMachine(verbose=self.verbose, networking=network.host(), image=BASE_IMAGE)
        self.machine_avocado.start()
        self.machine_avocado.wait_boot()
        self.machine_avocado.set_address("%s/20" % self.ipaddr_avocado)
        if not self.machine_cockpit:
            self.machine_cockpit = self.machine_avocado
        self._upload_library()
        self._cockpit_prepare()
        self.selenium_grid_restarts = []


    def _cockpit_prepare(self):
        self.machine_cockpit.wait_boot()
        self.machine_cockpit.set_address("%s/20" % self.ipaddr_cockpit)
        self.machine_cockpit.start_cockpit()
        self.machine_cockpit.execute("systemctl enable cockpit.socket")
        self._add_test_user()

    def cleanup(self):
        if self.machine_avocado:
            self.machine_avocado.kill()


    def _upload_tests(self, test_list, relative_path=CURRENT_DIR):
        self.machine_avocado.execute("mkdir -p {}".format(self.target_test_dir))
        test_list_expanded = []
        for test_item in test_list:
            test_list_expanded += glob.glob1(relative_path, test_item)
        self.machine_avocado.upload(test_list_expanded, self.target_test_dir, relative_dir=relative_path)

    def _add_test_user(self):
        self.machine_cockpit.execute("adduser test")
        self.machine_cockpit.execute("echo superhardpasswordtest5554 | passwd --stdin test")
        self.machine_cockpit.execute("usermod -a -G wheel test")
        self.machine_cockpit.execute("echo 'test        ALL=(ALL)       NOPASSWD: ALL' >> /etc/sudoers")
        self.machine_cockpit.execute("sed -i 's/^Defaults.*requiretty/#&/g' /etc/sudoers")
        self.machine_cockpit.execute("echo 'Defaults !requiretty' >> /etc/sudoers")

    def _upload_library(self):
        self.machine_avocado.execute("mkdir -p {}".format(self.target_library_path))
        self.machine_avocado.upload([os.path.join(CURRENT_DIR, self.library_name)], self.target_library_path)

    def _avocado_test_list(self, test_files):
        test_abspaths = [os.path.join(self.target_test_dir, x) for x in test_files]
        cmd_output = self.machine_avocado.execute(self.base_cmd + " list " + " ".join(test_abspaths))
        output_list = []
        for line in cmd_output.splitlines():
            output_list.append(line.split(" ", 1)[1].strip())
        return output_list


    def _run_one_avocado_test(self, test):
        status = True
        start_time = datetime.datetime.now()
        # wrap command to timeout to see also partial run result, to see what happens
        # instead of empty output and just timeout
        self.avocado_command = ["timeout {}".format(self.timeout - round(self.timeout/100)),
                     self.base_cmd, "run",
                     "--show-job-log",
                     test,
                     "2>&1",
                     ]
        try:
            cmd_output = self.machine_avocado.execute(command=" ".join(self.avocado_command),
                                         timeout=self.timeout,
                                         environment=self.env
                                         )
        except subprocess.CalledProcessError as exc:
            cmd_output = exc.output
            status = False
        except RuntimeError as exc:
            cmd_output = str(exc)
            status = False
        end_time = datetime.datetime.now()
        duration = (end_time - start_time).seconds

        return status, cmd_output, duration

    def run(self, test_list, timeout=3600):
        self.timeout = timeout
        self._upload_tests(test_list)
        tests = self._avocado_test_list(test_list)
        self.tap = TAPformatter(len(tests))
        start_time = datetime.datetime.now()
        self.tap.tap_header()
        for test in tests:
            self.tap.tap_item_header(test)
            status, cmd_output, duration = self._run_one_avocado_test(test)
            self.tap.tap_item_summary(test, cmd_output, duration)
            if not status:
                self.copy_artifacts()
        end_time = datetime.datetime.now()
        self.tap.tap_summary((end_time - start_time).seconds, self.selenium_grid_restarts)
        return self.tap.get_status()

    def copy_artifacts(self):
        if self.machine_avocado and self.machine_avocado.ssh_reachable:

            # get the screenshots first and move them on the guest so they won't get archived
            remote_screenshots = "*.png*"
            test_command_screenshots = "ls {}  && echo exists || echo No screenshot file found".format(remote_screenshots)
            if "exists" in self.machine_avocado.execute(command=test_command_screenshots):
                self.machine_avocado.download(remote_screenshots, arg_attachments_dir)

class SeleniumTestSuite(AvocadoTestSuite):
    machine_selenium = None
    ipaddr_selenium = "10.111.112.10"
    error_messages = [["urlopen error [Errno 113] No route to host"],
                      ["Timeout: Unable to attach remote Browser on hub"],
                      ["Terminated", "Running 'kill"]
                      ]

    def __init__(self, browser, **kwargs):
        super().__init__(**kwargs)
        self.env["HUB"] = self.ipaddr_selenium
        self.env["GUEST"] = self.ipaddr_cockpit
        self.env["BROWSER"] = browser
        self._prepare_selenium()
        self.wait_for_selenium_running(self.ipaddr_selenium)

    def _cockpit_prepare(self):
        super()._cockpit_prepare()
        self._prepare_machines_test_env()

    def cleanup(self):
        super().cleanup()
        if self.machine_selenium:
            self.machine_selenium.kill()

    def wait_for_selenium_running(self, host, port=4444):
        WAIT_SELENIUM_RUNNING = """#!/bin/sh
    until curl -s --connect-timeout 3 http://%s:%d >/dev/null; do
    sleep 0.5;
    done;
    """ % (host, port)
        with testvm.Timeout(seconds=300, error_message="Timeout while waiting for selenium to start"):
            self.machine_avocado.execute(script=WAIT_SELENIUM_RUNNING)

    def _prepare_machines_test_env(self):
        # Ensure everything has started correctly
        self.machine_cockpit.execute("systemctl start libvirtd.service")
        # Wait until we can get a list of domains
        self.machine_cockpit.execute("until virsh list > /dev/null; do sleep 0.5; done")
        # Wait for the network 'default' to become active
        self.machine_cockpit.execute("until virsh net-info default | grep Active > /dev/null; do sleep 0.5; done")
        self.machine_cockpit.execute("systemctl try-restart libvirtd")

        # Prepare image
        image_file = self.machine_cockpit.pull("cirros")
        self.machine_cockpit.upload([image_file], "/var/lib/libvirt/images")
        self.machine_cockpit.execute("cd /var/lib/libvirt/images && mv {} cirros.qcow2 && chmod 644 cirros.qcow2".format(
            os.path.basename(image_file)))

    def _prepare_selenium(self):
        if not self.machine_selenium:
            self.machine_selenium = testvm.VirtMachine(image="services", verbose=self.verbose, networking=self.network.host())
            # actually wait here, because starting selenium takes a while
            self.machine_selenium.pull(self.machine_selenium.image_file)
        self.machine_selenium.start()
        self.machine_selenium.wait_boot()
        self.machine_selenium.set_address("%s/20" % self.ipaddr_selenium)
        # start selenium on the server
        self.machine_selenium.upload(["selenium/selenium_start.sh"], "/root")
        self.machine_selenium.execute(command="/root/selenium_start.sh")

    def _restart_grid_and_cockpit(self):
        own_print("RESTARTING selenium grid and run test again")
        own_print("-------------------------------------------")
        # restart also cockpit machine to have clean env, because
        # some test could break the machine
        self.machine_selenium.kill()
        self.machine_cockpit.kill()
        self.machine_cockpit.start()
        self._prepare_selenium()
        self.wait_for_selenium_running(self.ipaddr_selenium)
        self._cockpit_prepare()

    def _run_one_avocado_test(self, test):
        status, cmd_output, duration = super()._run_one_avocado_test(test)
        if not status:
            # restart test again if it lost selenium grid location
            for error in self.error_messages:
                if all([x in cmd_output for x in error]):
                    self._restart_grid_and_cockpit()
                    self.selenium_grid_restarts.append(test)
                    status, cmd_output, duration = super()._run_one_avocado_test(test)
                    break
        return status, cmd_output, duration


class SeleniumWinTestSuite(SeleniumTestSuite):
    def _prepare_selenium(self):
        if not self.machine_selenium:
            self.machine_avocado.dhcp_server(range=[self.ipaddr_selenium, self.ipaddr_selenium])
            self.machine_selenium = testvm.VirtMachine(image='windows-10', verbose=self.verbose, networking=self.network.host(), memory_mb=2048, cpus=2)
            self.machine_selenium.pull(self.machine_selenium.image_file)
        self.machine_selenium.start()


class SeleniumOwnHub(SeleniumTestSuite):
    def __init__(self, address_string, browser, **kwargs):
        super().__init__(**kwargs)
        self.browser = browser
        address_option_list = address_string.split(":")
        hub_address = address_option_list[0]
        if len(address_option_list) > 1:
            guest_adress = address_option_list[1]
        else:
            guest_adress = "127.0.0.2"
        if len(address_option_list) > 2:
            cockpit_port = address_option_list[2]
        else:
            cockpit_port = self.machine_cockpit.forward["9090"].split(":")[1]

        self.env["HUB"] = hub_address
        self.env["GUEST"] = guest_adress
        self.env["PORT"] = cockpit_port
        # avocado and cockpit machine is in same network so use internal address
        self.env["SSH_PORT"] = "22"
        self.env["SSH_GUEST"] = self.ipaddr_cockpit
        self.env["BROWSER"] = browser
        self.wait_for_selenium_running(hub_address)

    def _restart_grid_and_cockpit(self):
        own_print("RESTARTING cockpit machine")
        own_print("-------------------------------------------")
        # restart also cockpit machine to have clean env, because
        # some test could break the machine
        self.machine_cockpit.kill()
        self.machine_cockpit.start()
        self._cockpit_prepare()


def main():
    parser = argparse.ArgumentParser(description='Run Cockpit Avocado test(s)')
    parser.add_argument('-v', '--verbose', dest="verbosity", action='store_true',
                        help='Verbose output')
    parser.add_argument("--hub",
                        dest="hub",
                        action="store",
                        help="Use own Selenium hub (it is important to ensure that hub can see cockpit machine)"
                             " You can use address formats:"
                             " yourHubIP"
                             " or yourHubIP:CockpitIP"
                             " or yourHubIP:CockpitIP:CockpitPort"
                        )
    parser.add_argument("-b", "--browser", choices=['firefox', 'chrome', 'edge'],
                    default='firefox',
                    help="selenium browser choice (default: %(default)s)")
    parser.add_argument("-l", "--logs", dest='download_logs', action='store_true',
                        help="Always download avocado logs, even on success")
    parser.add_argument("--sit", dest='sit', action='store_true',
                        help="Sit and wait after test failure")
    parser.add_argument('tests', nargs='*', default=["selenium-*.py"])

    opts = parser.parse_args()

    image = BASE_IMAGE
    # When Windows in use, we need rtl8139, for anything else we want virtio-net-pci
    network = testvm.VirtNetwork(image="windows-10" if "edge" in opts.browser else image)
    machine = None
    test_timeout = TEST_TIMEOUT
    testsuite = None
    try:
        machine = testvm.VirtMachine(verbose=opts.verbosity, networking=network.host(), image=image)
        machine.start()
        if opts.hub:
            testsuite = SeleniumOwnHub(address_string=opts.hub,
                                       browser=opts.browser,
                                       machine_cockpit=machine,
                                       network=network,
                                       verbose=opts.verbosity)
        elif 'edge' in opts.browser:
            testsuite = SeleniumWinTestSuite(browser=opts.browser, machine_cockpit=machine, network=network, verbose=opts.verbosity)
            test_timeout = TEST_TIMEOUT * 2
        else:
            testsuite = SeleniumTestSuite(browser=opts.browser, machine_cockpit=machine, network=network, verbose=opts.verbosity)

        success = testsuite.run(opts.tests, timeout=test_timeout)
        if not success and opts.sit:
            avocado_command = " ".join(["{}={}".format(k, v) for k, v in testsuite.env.items()]
                                       + testsuite.avocado_command)
            sys.stderr.write(testsuite.machine_cockpit.diagnose() +
                             "\nTo rerun tests connect to VM with avocado installed via" +
                             "\n\tssh -p {} root@127.0.0.2\n".format(
                                 testsuite.machine_avocado.forward["22"].split(":")[1]) +
                             "\nand run the command" +
                             "\n\t " + avocado_command +
                             "\n\nPress RET to continue...\n"
                             )
            sys.stdin.readline()
    finally:
        if machine:
            machine.kill()
        if testsuite:
            testsuite.cleanup()

    if success:
        return 0
    else:
        return 1


if __name__ == '__main__':
    sys.exit(main())