#!/usr/bin/python3
# This can also be run with Python 2.

# Copyright (C) 2013 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/>.

#
# This is a TAP driver for automake
#
# In particular it leaves stderr untouched, and is cleaner than the
# one implemented in shell that is making the rounds.
#
# This implements the automake "Custom Test Driver" protocol:
# https://www.gnu.org/software/automake/manual/html_node/Custom-Test-Drivers.html
#
# This consumes the Test Anything Protocol (ie: TAP)
# https://metacpan.org/pod/release/PETDANCE/Test-Harness-2.64/lib/Test/Harness/TAP.pod
#

import argparse
import fcntl
import os
import select
import struct
import subprocess
import sys
import termios
import errno

_PY3 = sys.version[0] >= '3'
_str = _PY3 and str or unicode

def out(data, stream=None, flush=False):
    if not isinstance(data, bytes):
        data = data.encode("UTF-8")
    if not stream:
        stream = _PY3 and sys.stdout.buffer or sys.stdout
    while True:
        try:
            if data:
                stream.write(data)
            data = None
            if flush:
                stream.flush()
            flush = False
            break
        except IOError as e:
            if e.errno == errno.EAGAIN:
                continue
            raise

def terminal_width():
    try:
        h, w, hp, wp = struct.unpack('HHHH',
            fcntl.ioctl(1, termios.TIOCGWINSZ,
            struct.pack('HHHH', 0, 0, 0, 0)))
        return w
    except IOError as e:
        if e.errno != errno.ENOTTY:
            sys.stderr.write("%i %s %s\n" % (e.errno, e.strerror, sys.exc_info()))
        return sys.maxsize

class Driver:
    def __init__(self, args):
        self.argv = args.command
        self.test_name = args.test_name
        self.log = open(args.log_file, "wb", 0)
        self.log.write(("# %s\n" % " ".join(sys.argv)).encode("UTF-8"))
        self.trs = open(args.trs_file, "w", 1)
        self.color_tests = args.color_tests
        self.expect_failure = args.expect_failure
        self.enable_hard_errors = args.enable_hard_errors
        self.width = terminal_width() - 9

    def report(self, code, *args):
        CODES = {
            "XPASS": '\x1b[0;31m', # red
            "FAIL": '\x1b[0;31m', # red
            "PASS": '\x1b[0;32m', # grn
            "XFAIL": '\x1b[1;32m', # lgn
            "SKIP": '\x1b[1;34m', # blu
            "ERROR": '\x1b[0;35m', # mgn
        }

        # Print out to console
        if self.color_tests:
            if code in CODES:
                out(CODES[code])
        out(code)
        if self.color_tests:
            out('\x1b[m')
        out(": ")
        msg = "".join([ self.test_name + " " ] + list(map(_str, args)))
        if code == "PASS" and len(msg) > self.width:
            out(msg[:self.width])
            out("...")
        else:
            out(msg)
        out("\n", flush=True)

        # Book keeping
        if code in CODES:
            self.trs.write(":test-result: %s\n" % code)

    def result_pass(self, *args):
        if self.expect_failure:
            self.report("XPASS", *args)
        else:
            self.report("PASS", *args)

    def result_fail(self, *args):
        if self.expect_failure:
            self.report("XFAIL", *args)
        else:
            self.report("FAIL", *args)

    def result_skip(self, *args):
        if self.expect_failure:
            self.report("XFAIL", *args)
        else:
            self.report("SKIP", *args)

    def report_error(self, description=""):
        if self.enable_hard_errors:
            self.report("ERROR", "", description)
        else:
            self.result_fail(description)

    def process(self, output):
        pass

    def execute(self):
        try:
            proc = subprocess.Popen(self.argv, close_fds=True,
                                    stdin=subprocess.PIPE,
                                    stdout=subprocess.PIPE,
                                    stderr=subprocess.PIPE)
        except OSError as ex:
            self.report_error("Couldn't run %s: %s" % (self.argv[0], str(ex)))
            return

        proc.stdin.close()
        outf = proc.stdout.fileno()
        errf = proc.stderr.fileno()
        rset = [outf, errf]
        while len(rset) > 0:
            ret = select.select(rset, [], [], 10)
            if outf in ret[0]:
                data = os.read(outf, 1024)
                if data == b"":
                    rset.remove(outf)
                self.log.write(data)
                self.process(data)
            if errf in ret[0]:
                data = os.read(errf, 1024)
                if data == b"":
                    rset.remove(errf)
                self.log.write(data)
                stream = _PY3 and sys.stderr.buffer or sys.stderr
                out(data, stream=stream, flush=True)

        proc.wait()

        # Make sure the test didn't change blocking output
        assert fcntl.fcntl(0, fcntl.F_GETFL) & os.O_NONBLOCK == 0
        assert fcntl.fcntl(1, fcntl.F_GETFL) & os.O_NONBLOCK == 0
        assert fcntl.fcntl(2, fcntl.F_GETFL) & os.O_NONBLOCK == 0

        return proc.returncode


class TapDriver(Driver):
    def __init__(self, args):
        Driver.__init__(self, args)
        self.output = ""
        self.reported = { }
        self.test_plan = None
        self.late_plan = False
        self.errored = False
        self.bail_out = False
        self.skip_all_reason = None

    def report(self, code, num, *args):
        if num:
            Driver.report(self, code, num, " ", *args)
            self.reported[num] = code
        else:
            Driver.report(self, code, *args)
        if code == "ERROR":
            self.errored = True

    def consume_test_line(self, ok, data):
        # It's an error if the caller sends a test plan in the middle of tests
        if self.late_plan:
            self.report_error("Got tests after late TAP test plan")
            self.late_plan = False

        # Parse out a number and then description
        (num, unused, description) = data.partition(" ")
        try:
            num = int(num)
        except ValueError:
            self.report_error("Invalid test number: %s" % data)
            return
        description = description.lstrip()

        # Parse out a directive from description, if any
        (description, unused, directive) = description.partition("#")

        # Special case if directive starts with this, then skip
        if directive.lstrip().lower().startswith("skip"):
            self.result_skip(num, description)
        elif ok:
            self.result_pass(num, description)
        else:
            self.result_fail(num, description)

    def consume_test_plan(self, line):
        # Only one test plan is supported
        if self.test_plan:
            self.report_error("Get a second TAP test plan")
            return

        if line.lower().startswith('1..0 # skip'):
            self.skip_all_reason = line[5:].strip()
            self.bail_out = True
            return

        try:
            (first, unused, last) = line.partition("..")
            first = int(first)
            last = int(last)
        except ValueError:
            self.report_error("Invalid test plan: %s..%s" % (first, last))
            return

        self.test_plan = (first, last)
        self.late_plan = self.reported and True or False

    def consume_bail_out(self, line):
        self.bail_out = True
        self.report("SKIP", 0, line)

    def process(self, output):
        if output:
            self.output += output.decode("UTF-8")
        elif self.output:
            self.output += "\n"
        (ready, unused, self.output) = self.output.rpartition("\n")
        for line in ready.split("\n"):
            if line.startswith("ok "):
                self.consume_test_line(True, line[3:])
            elif line.startswith("not ok "):
                self.consume_test_line(False, line[7:])
            elif line and line[0].isdigit() and ".." in line:
                self.consume_test_plan(line)
            elif line.lower().startswith("bail out!"):
                self.consume_bail_out(line)

    def run(self):
        returncode = self.execute()

        failed = False
        skipped = True

        if self.skip_all_reason is not None:
            self.result_skip("skipping:", self.skip_all_reason)
            self.trs.write(":global-test-result: SKIP\n")
            self.trs.write(":test-global-result: SKIP\n")
            self.trs.write(":recheck: no\n")
            return 0

        # Basic collation of results
        for (num, code) in self.reported.items():
            if code == "ERROR":
                self.errored = True
            elif code == "FAIL" or code == "XPASS":
                failed = True
            if code != "SKIP":
                skipped = False

        if not self.errored:
            if returncode == 77:
                skipped = True
            elif returncode:
                self.report_error("process failed: %d" % returncode)
                self.errored = True

        # Check the plan
        if not self.errored:
            if not self.test_plan:
                if not self.bail_out:
                    self.report_error("Didn't receive a TAP test plan")
            else:
                for i in range(self.test_plan[0], self.test_plan[1] + 1):
                    if i not in self.reported:
                        if self.bail_out:
                            self.report("SKIP", i, "- bailed out")
                        else:
                            self.report("ERROR", i, "- missing test")
                            skipped = False
                            self.errored = True

        if self.errored:
            self.trs.write(":global-test-result: ERROR\n")
            self.trs.write(":test-global-result: ERROR\n")
            self.trs.write(":recheck: yes\n")
        elif failed:
            self.trs.write(":global-test-result: FAIL\n")
            self.trs.write(":test-global-result: FAIL\n")
            self.trs.write(":recheck: yes\n")
        elif skipped:
            self.trs.write(":global-test-result: SKIP\n")
            self.trs.write(":test-global-result: SKIP\n")
            self.trs.write(":recheck: no\n")
        else:
            self.trs.write(":global-test-result: PASS\n")
            self.trs.write(":test-global-result: PASS\n")
            self.trs.write(":recheck: no\n")
        if self.errored or failed:
            self.trs.write(":copy-in-global-log: yes\n")

        # Process result code
        return 0


class SimpleDriver(Driver):
    def __init__(self, args):
        Driver.__init__(self, args)

    def run(self):
        returncode = self.execute()
        if returncode == 0:
            self.result_pass()
            self.trs.write(":global-test-result: PASS\n")
            self.trs.write(":test-global-result: PASS\n")
            self.trs.write(":recheck: no\n")
        elif returncode == 77:
            self.result_skip()
            self.trs.write(":global-test-result: SKIP\n")
            self.trs.write(":test-global-result: SKIP\n")
            self.trs.write(":recheck: no\n")
        elif returncode == 99:
            self.report_error()
            self.trs.write(":global-test-result: ERROR\n")
            self.trs.write(":test-global-result: ERROR\n")
            self.trs.write(":copy-in-global-log: yes\n")
            self.trs.write(":recheck: yes\n")
        else:
            self.result_fail()
            self.trs.write(":global-test-result: FAIL\n")
            self.trs.write(":test-global-result: FAIL\n")
            self.trs.write(":copy-in-global-log: yes\n")
            self.trs.write(":recheck: yes\n")

        # Process result code
        return 0


class MissingDriver(Driver):
    def __init__(self, args):
        Driver.__init__(self, args)
        self.missing = args.missing

    def run(self):
        self.result_skip("skipping due to: ", self.missing)
        self.trs.write(":global-test-result: SKIP\n")
        self.trs.write(":test-global-result: SKIP\n")
        self.trs.write(":recheck: no\n")
        return 0


class YesNoAction(argparse.Action):
    def __init__(self, option_strings, dest, **kwargs):
        argparse.Action.__init__(self, option_strings, dest, **kwargs)
        self.metavar = "[yes|no]"
    def __call__(self, parser, namespace, values, option_string=None):
        if not values or "yes" in values:
            setattr(namespace, self.dest, True)
        else:
            setattr(namespace, self.dest, False)


def main(argv):
    parser = argparse.ArgumentParser(description='Automake TAP driver')
    parser.add_argument('--format', metavar='FORMAT', choices=[ "simple", "tap" ],
                        default="tap", help='The type of test to drive')
    parser.add_argument('--missing', metavar="TOOL", nargs='?',
                        help="Force the test to skip due to missing tool")
    parser.add_argument('--test-name', metavar='NAME',
                        help='The name of the test')
    parser.add_argument('--log-file', metavar='PATH.log', required=True,
                        help='The .log file the driver creates')
    parser.add_argument('--trs-file', metavar='PATH.trs', required=True,
                        help='The .trs file the driver creates')
    parser.add_argument('--color-tests', default=True, action=YesNoAction,
                        help='Whether the console output should be colorized or not')
    parser.add_argument('--expect-failure', default=False, action=YesNoAction,
                        help="Whether the tested program is expected to fail")
    parser.add_argument('--enable-hard-errors', default=False, action=YesNoAction,
                        help="Whether hard errors in the tested program are treated differently")
    parser.add_argument('command', nargs='+',
                        help="A test command line to run")
    args = parser.parse_args(argv[1:])

    if not args.test_name:
        args.test_name = os.path.basename(args.command[0])
    if args.missing:
        driver = MissingDriver(args)
    elif args.format == "simple":
        driver = SimpleDriver(args)
    elif args.format == "tap":
        driver = TapDriver(args)
    return driver.run()

if __name__ == "__main__":
    sys.exit(main(sys.argv))