#!/usr/bin/env python3

import argparse
import glob
import os
import sys

from subprocess import PIPE, DEVNULL, STDOUT


class Runcmd:
    """Run external commands in various ways"""

    def __init__(self, verbose, progress):
        self._verbose = verbose
        self._progress = progress

    def _write_msg(self, msg):
        sys.stdout.write(f"{msg}\n")
        sys.stdout.flush()

    def msg(self, msg):
        if self._verbose:
            self._write_msg(msg)

    def title(self, title):
        if self._verbose:
            self._write_msg("")
            self._write_msg("=" * len(title))
            self._write_msg(title)
            self._write_msg("=" * len(title))
        elif self._progress:
            self._write_msg(title)

    def _runcmd_unchecked(self, argv, **kwargs):
        """Run a command (generic version)

        Return a subcommand.CompletedProcess. It's the caller's duty
        check that the command succeeded.

        All actual execution of other programs happens via this method.
        However, only methods of this class should ever call this
        method.
        """

        # Import "run" here so that no other part of the code can see the
        # symbol.
        from subprocess import run

        self.msg(f"RUN: {argv} {kwargs}")

        if not self._verbose:
            kwargs["stdout"] = PIPE
            kwargs["stderr"] = STDOUT

        return run(argv, **kwargs)

    def runcmd(self, argv, **kwargs):
        """Run a command (generic version)

        If the command fails, terminate the program. On success, return
        a subprocess.CompletedProcess.
        """

        p = self._runcmd_unchecked(argv, **kwargs)
        if p.returncode != 0:
            sys.stdout.write((p.stdout or b"").decode("UTF-8"))
            sys.stderr.write((p.stderr or b"").decode("UTF-8"))
            sys.stderr.write(f"Command {argv} failed\n")
            sys.exit(p.returncode)
        return p

    def runcmd_maybe(self, argv, **kwargs):
        """Run a command it's availbe, or quietly do no nothing"""
        if self.got_command(argv[0]):
            self.runcmd(argv, **kwargs)

    def got_command(self, name):
        """Is a command of a given name available?"""
        p = self._runcmd_unchecked(["which", name], stdout=DEVNULL)
        return p.returncode == 0

    def cargo(self, args, **kwargs):
        """Run cargo with arguments."""
        self.runcmd(["cargo"] + args, **kwargs)

    def cargo_maybe(self, args, **kwargs):
        """Run cargo if the desired subcommand is available"""
        if self.got_cargo(args[0]):
            self.runcmd(["cargo"] + args, **kwargs)

    def got_cargo(self, subcommand):
        """Is a cargo subcommand available?"""
        p = self.runcmd(["cargo", "--list"], check=True, stdout=PIPE)
        lines = [line.strip() for line in p.stdout.decode("UTF-8").splitlines()]
        return subcommand in lines

    def codegen(self, md, output, **kwargs):
        """Run the Subplot code generator and the test program it produces"""
        self.cargo(
            [
                "run",
                "--package=subplot",
                "--bin=subplot",
                "--",
                f"--resources={os.path.abspath('share')}",
                "codegen",
                md,
                f"--output={output}",
            ],
            **kwargs,
        )

    def docgen(self, md, output, **kwargs):
        """Run the Subplot document generator"""
        self.cargo(
            [
                "run",
                "--package=subplot",
                "--bin=subplot",
                "--",
                f"--resources={os.path.abspath('share')}",
                "docgen",
                md,
                f"--output={output}",
            ],
            **kwargs,
        )


def find_files(pattern, pred):
    """Find files recursively, if they are accepted by a predicate function"""
    return [f for f in glob.glob(pattern, recursive=True) if pred(f)]


def check_python(r):
    """Run all checks for Python code"""
    r.title("checking Python code")

    # Find all Python files anywhere, except those we know aren't proper Python.
    py = find_files(
        "**/*.py", lambda f: os.path.basename(f) not in ("template.py", "test.py")
    )

    # Test with flake8 if available. Flake8 finds Python files itself.
    r.runcmd_maybe(["flake8", "check"] + py)

    # Check formatting with Black. We need to provide the files to Python
    # ourselves.
    r.runcmd_maybe(["black", "--check"] + py)

    # Find and run unit tests.
    tests = find_files("**/*_tests.py", lambda f: True)
    for test in tests:
        dirname = os.path.dirname(test)
        test = os.path.basename(test)
        r.runcmd(["python3", test], cwd=dirname)


def check_shell(r):
    """Run all checks for shell code"""
    r.title("checking shell code")

    # Find all shell files anywhere, except generated test programs.
    sh = find_files("**/*.sh", lambda f: os.path.basename(f) != "test.sh")

    r.runcmd_maybe(["shellcheck"] + sh)


def check_rust(r, strict=False):
    """Run all checks for Rust code"""
    r.title("checking Rust code")

    r.runcmd(["cargo", "build", "--all-targets"])
    if r.got_cargo("clippy"):
        if strict:
            r.runcmd(["cargo", "clippy", "--all-targets", "--", "-D", "warnings"])
        else:
            r.runcmd(["cargo", "clippy", "--all-targets"])
    elif strict:
        sys.exit("Strict Rust checks specified, but clippy was not found")
    r.runcmd(["cargo", "test"])
    r.runcmd(["cargo", "fmt"])


def check_subplots(r):
    """Run all Subplots and generate documents for them"""
    mds = find_files("**/*.md", lambda f: f == f.lower() and "subplotlib" not in f)
    for md0 in mds:
        r.title(f"checking subplot {md0}")

        dirname = os.path.dirname(md0) or "."
        md = os.path.basename(md0)

        template = get_template(md0)
        if template == "python":
            # Remove test log from previous run, if any.
            test_log = os.path.join(dirname, "test.log")
            if os.path.exists(test_log):
                os.remove(test_log)

            r.codegen(md, "test.py", cwd=dirname)
            r.runcmd(["python3", "test.py", "--log", "test.log"], cwd=dirname)
            os.remove(test_log)
            os.remove(os.path.join(dirname, "test.py"))
        elif template == "bash":
            r.codegen(md, "test.sh", cwd=dirname)
            r.runcmd(["bash", "-x", "test.sh"], cwd=dirname)
            os.remove(os.path.join(dirname, "test.sh"))
        else:
            sys.exit(f"unknown template {template} in {md0}")

        base, _ = os.path.splitext(md)
        r.docgen(md, base + ".pdf", cwd=dirname)
        r.docgen(md, base + ".html", cwd=dirname)


def check_subplotlib(r):
    """Run all checks for subplotlib"""
    r.title("checking subplotlib code")

    # Run Rust tests for the subplotlib library.
    r.runcmd(["cargo", "test", "--lib"], cwd="subplotlib")

    # Find subplots for subplotlib.
    mds = find_files("subplotlib/*.md", lambda f: True)

    os.makedirs("subplotlib/tests", exist_ok=True)
    for md0 in mds:
        r.title(f"checking subplot {md0}")
        dirname = os.path.dirname(md0)
        md = os.path.basename(md0)
        base, _ = os.path.splitext(md)
        test_rs = os.path.join("tests", base + ".rs")
        r.codegen(md, test_rs, cwd=dirname)
        r.cargo(["fmt"], cwd=dirname)
        r.cargo(["test", "--test", base], cwd=dirname)
        r.docgen(md, base + ".html", cwd=dirname)
        r.docgen(md, base + ".pdf", cwd=dirname)


def get_template(filename):
    prefix = "template: "
    with open(filename) as f:
        data = f.read()
        for line in data.splitlines():
            if line.startswith(prefix):
                line = line[len(prefix) :]
                return line
    sys.exit(f"{filename} does not specify a template")


def parse_args():
    """Parse command line arguments to this script"""
    p = argparse.ArgumentParser()
    p.add_argument("-v", dest="verbose", action="store_true", help="be verbose")
    p.add_argument(
        "-p", dest="progress", action="store_true", help="print some progress output"
    )
    p.add_argument(
        "--strict", action="store_true", help="don't allow compiler warnings"
    )

    all_whats = ["python", "shell", "rust", "subplots", "subplotlib"]
    p.add_argument(
        "what", nargs="*", default=all_whats, help=f"what to test: {all_whats}"
    )
    return p.parse_args()


def main():
    """Main program"""
    args = parse_args()

    r = Runcmd(args.verbose, args.progress)

    for what in args.what:
        if what == "python":
            check_python(r)
        elif what == "shell":
            check_shell(r)
        elif what == "rust":
            check_rust(r, strict=args.strict)
        elif what == "subplots":
            check_subplots(r)
        elif what == "subplotlib":
            check_subplotlib(r)
        else:
            sys.exit(f"Unknown test {what}")

    sys.stdout.write("Everything seems to be in order.\n")


main()