#!/usr/bin/env python3

import os
import sys
import subprocess
import re
import multiprocessing
from concurrent.futures import ThreadPoolExecutor
from common import *

usage="""Usage:
    rmake --all                  Compile all configured suites

    rmake --list                 Show available compilation targets
    rmake --doc                  Build & serve documentation (rebuilds on every change)
    rmake --help                 Print this help

    rmake <targetName>...        Compile one or multiple targets
    rmake <pathPrefix>...        Compile one or multiple targets in path
    rmake <suitename>...         Compile given suite

    rmake --tidy <targetName>... Run clang-tidy on one or multiple targets
    rmake --tidy <pathPrefix>... Run clang-tidy on one or multiple targets in path
    rmake --tidy --all           Run clang-tidy on all configured targets

    rmake --test --all               Compile and run all configured tests
    rmake --test <testTargetName>... Compile and run test for target
    rmake --test <pathName>...       Compile and run tests from given path prefix
    rmake --test <suitename>...      Compile and run tests from given suite

    rmake clean                  Clean all suite builds
    rmake clean-<suitename>...   Clean given suite
"""

def invokeBuild(wd, targets):
    """
    Invoke ninja or make based on the environment for targets in a working
    directory
    """
    maxThreads = 2 + multiprocessing.cpu_count()
    command = ["cmake", "--build", wd, f"-j{maxThreads}"]
    if len(targets) > 0:
        command += ["--target"] + targets
    subprocess.check_call(command, cwd=wd)

def compileTargets(targets):
    """
    Compile given targets. Be as parallel as possible
    """
    for s in configuredSuites():
        sTargets = [x.target for x in targets if x.suite == s]
        if len(sTargets) == 0:
            continue
        print(f"Compiling targets from suite {s}: {', '.join(sTargets)}")
        try:
            invokeBuild(os.path.join(buildDir, s), sTargets)
        except Exception as e:
            sys.exit(f"Build of target {', '.join(sTargets)} failed with {e}. See output above\n")

def compileSuite(suite):
    print(f"Compiling suite {suite}")
    try:
        invokeBuild(os.path.join(buildDir, suite), [])
    except Exception as e:
        sys.exit(f"Build of suite {suite} failed with {e}. See output above\n")

def cleanSuite(suite):
    print(f"Cleaning suite {suite}")
    try:
        wd = os.path.join(buildDir, suite)
        subprocess.check_call(["cmake", "--build", wd, "--target", "clean"],
            cwd=wd)
    except subprocess.SubprocessError as e:
        sys.exit(f"Cleaning of suite {suite} failed. See output above\n")

def runTest(target):
    testname = f"{target.suite}::{target.target}"
    print(f"Running test: {testname}")
    try:
        subprocess.check_call([target.target])
    except subprocess.SubprocessError as e:
        sys.exit(f"Running of suite {testname} failed. See output above\n")

def printTargets():
    aSuites = availableSuites()
    aTargets = availableTargets()
    print("Available suites: " + ", ".join(aSuites) + "\n")
    for suite in aSuites:
        print(f"Available targets for suite {suite}:")
        empty = True
        for target in aTargets:
            if target.suite != suite:
                continue
            empty = False
            print(f"  {target.target}: {target.path}")
        if empty:
            print("  None")
        print("")

def allAreTargets(args):
    for a in args:
        if a.startswith("--"):
            return False
    return True

def expandSuite(suite, targetList):
    return [x for x in targetList if x.suite == suite]

def matchPattern(pattern, allowedSuites, allowedTargets, suitesToCompile,
    targetsToCompile, sourcesToCompile, suitesToClean):
    """
    Match pattern to allowed suites and targets and add to either
    suitesToCompile, targetsToCompile or suitesToClean
    """
    if os.path.exists(pattern) and os.path.isfile(pattern):
        sourcesToCompile.add(pattern)
        return
    if pattern in allowedSuites:
        suitesToCompile.add(pattern)
        return
    used = False
    for target in allowedTargets:
        if target.path.startswith(pattern) or target.target == pattern:
            targetsToCompile.add(target)
            used = True
    if pattern == "clean":
        suitesToClean.update(allowedSuites)
        used = True
    if pattern.startswith("clean-") and pattern[6:] in allowedSuites:
        suitesToClean.add(pattern[6:])
        used = True
    if not used:
        sys.exit(f"Invalid target pattern specified: '{pattern}'")

def build(patterns):
    if not allAreTargets(patterns):
        sys.exit("Invalid usage!\n\n" + usage)

    aSuites = configuredSuites()
    aTargets = availableTargets()

    suitesToCompile = set()
    targetsToCompile = set()
    sourcesToCompile = set()
    suitesToClean = set()

    for pattern in patterns:
        matchPattern(pattern, aSuites, aTargets, suitesToCompile,
            targetsToCompile, sourcesToCompile, suitesToClean)
    if len(sourcesToCompile) > 0:
        sys.exit("Cannot specify files for test")
    targetsToCompile = list(targetsToCompile)
    targetsToCompile.sort(key=lambda x: x.suite)

    for s in suitesToClean:
        cleanSuite(s)

    for s in suitesToCompile:
        compileSuite(s)
    compileTargets(targetsToCompile)

def buildAll():
    for suite in configuredSuites():
        if os.path.exists(os.path.join(buildDir, suite)):
            compileSuite(suite)
    sys.exit(0)

def test(patterns):
    aSuites = configuredSuites()
    aTargets = availableTestTargets()

    if len(patterns) == 0:
        sys.exit("Invalid usage!\n\n" + usage)
    if patterns[0] == "--all":
        if len(patterns) != 1:
            sys.exit("Invalid usage!\n\n" + usage)
        patterns = [x.target for x in aTargets]
    if not allAreTargets(patterns):
        sys.exit("Invalid usage!\n\n" + usage)

    suitesToCompile = set()
    targetsToCompile = set()
    sourcesToCompile = set()
    suitesToClean = set()
    for pattern in patterns:
        matchPattern(pattern, aSuites, aTargets, suitesToCompile,
            targetsToCompile, sourcesToCompile, suitesToClean)
    if len(sourcesToCompile) > 0:
        sys.exit("Cannot specify files for test")

    for s in suitesToCompile:
        targetsToCompile = targetsToCompile.union(expandSuite(s, aTargets))

    targetsToCompile = list(targetsToCompile)
    targetsToCompile.sort(key=lambda x: x.suite)

    for s in suitesToClean:
        cleanSuite(s)

    compileTargets(targetsToCompile)

    for t in targetsToCompile:
        runTest(t)

def suiteToSources(suite, targetList, sources):
    retSources = set()
    for target in targetList:
        if target.suite == suite:
            retSources = retSources.union(sources[suite][target.target])
    return retSources

def getClangTidyVersion():
    output = subprocess.check_output(["clang-tidy", "--version"]).decode("utf-8")
    match = re.search(r'LLVM version (\d+).(\d+).(\d+)', output)
    if match is not None:
        return int(match.group(1)), int(match.group(2)), int(match.group(3))
    raise RuntimeError(f"Cannot extract version from: {output}")

def runClangTidy(suite, source):
    """
    Run clang-tidy in a subprocess. Returns a tuple - success flag and stout
    """
    columns, _ = os.get_terminal_size()
    message = f"➡️  "
    path = f"{os.path.relpath(source)}"
    if len(path) + len(message) > columns:
        l = -columns + len(path) + len(message) + 4
        path = "..." + path[l:]
    GREEN = "\033[32m"
    ENDC = "\033[m"
    BOLD = "\u001b[1m"

    message += GREEN + BOLD + path + ENDC + "\n" + "―" * columns

    if not os.path.exists(source):
        message += "\nWarning: file does not exist, skipping."
        message += "\n   The source file is probably generated and it has not been generated yet\n"
        return True, message
    with open(os.path.join(suiteDir, suite, ".clang-tidy")) as f:
        configuration = f.read()
    major, _, _ = getClangTidyVersion()
    command = ["clang-tidy"]
    if major >= 12:
        command.append("--use-color")
    command += ["-p", os.path.join(buildDir, suite),
                "-header-filter=*",
                "--config", configuration,
                os.path.relpath(source)]

    proc = subprocess.run(command,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE)

    message += "\n" + proc.stdout.decode("utf-8") + "\n" + proc.stderr.decode("utf-8") + "\n"
    return proc.returncode == 0, message

def tidy(patterns):
    """
    Run clang-tidy on the given patterns
    """
    aSuites = configuredSuites()
    aTargets = availableTargets()
    aTestTargets = availableTestTargets()
    sources = sourceList()

    if len(patterns) == 0:
        sys.exit("Invalid usage!\n\n" + usage)
    if patterns[0] == "--all":
        if len(patterns) != 1:
            sys.exit("Invalid usage!\n\n" + usage)
        patterns = [x.target for x in aTargets]
    if not allAreTargets(patterns):
        sys.exit("Invalid usage!\n\n" + usage)

    suitesToCheck = set()
    targetsToCheck = set()
    sourcesToCheck = set()
    suitesToClean = set()
    for pattern in patterns:
        matchPattern(pattern, aSuites, aTargets, suitesToCheck,
            targetsToCheck, sourcesToCheck, suitesToClean)
    if len(suitesToClean):
        sys.exit("Cannot specify clean targets in clang-tidy check")

    success = True
    filesChecked = 0
    for currentSuite in aSuites:
        suiteSources = set()
        for x in sources[currentSuite].values():
            suiteSources = suiteSources.union(x)

        relevantSources = set()
        for s in sourcesToCheck:
            if os.path.realpath(s) in suiteSources:
                relevantSources.add(s)
        if currentSuite in suitesToCheck:
            for x in suiteToSources(currentSuite, aTargets, sources):
                relevantSources.add(x)
        else:
            for t in targetsToCheck:
                relevantSources = relevantSources.union(sources[currentSuite][t.target])

        if len(relevantSources) > 0:
            print(f"Running clang-tidy for suite: {currentSuite}\n")

        # Ignore tests
        ignoredSources = suiteToSources(currentSuite, aTestTargets, sources)
        relevantSources = relevantSources.difference(ignoredSources)

        # Make the order deterministic
        relevantSources = list(relevantSources)
        relevantSources.sort()
        # Skip build files
        relevantSources = [x for x in relevantSources if not x.startswith(buildDir)]

        # Run clang-tidy
        jobs = []
        for source in relevantSources:
            _, extension = os.path.splitext(source)
            if extension in [".c", ".cc", ".cpp"]:
                jobs.append(source)
        with ThreadPoolExecutor() as ex:
            for x in ex.map(lambda source: runClangTidy(currentSuite, source), jobs):
                valid, report = x
                success = success and valid
                filesChecked += 1
                print(report)
    if filesChecked == 0:
        print("No files to check")
        sys.exit(0)
    if success:
        print("✅ All checks passed")
        sys.exit(0)
    else:
        print("❌ Some checks failed")
        sys.exit(1)

def main():
    if len(sys.argv) < 2:
        sys.exit("Invalid usage!\n\n" + usage)

    command = sys.argv[1]

    if command == "--list":
        if len(sys.argv) != 2:
            sys.exit("Invalid usage!\n\n" + usage)
        printTargets()
        sys.exit(0)

    if command == "--doc":
        if len(sys.argv) != 2:
            sys.exit("Invalid usage!\n\n" + usage)
        try:
            tool = os.path.join(root, "releng", "tools", "_build_doc.sh")
            subprocess.check_call([tool])
            sys.exit(0)
        except subprocess.SubprocessError as e:
            sys.exit(1)
        except KeyboardInterrupt:
            sys.exit(0)

    if command == "--help":
        if len(sys.argv) != 2:
            sys.exit("Invalid usage!\n\n" + usage)
        print(usage)
        sys.exit(0)

    if command == "--test":
        test(sys.argv[2:])
        sys.exit(0)

    if command == "--tidy":
        tidy(sys.argv[2:])
        sys.exit(0)

    if command == "--all":
        if len(sys.argv) != 2:
            sys.exit("Invalid usage!\n\n" + usage)
        buildAll()
        sys.exit(0)

    build(sys.argv[1:])

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        sys.exit("Interrupted.")
