"""This is the main tool for handling llvm builds, bisects etc."""

import os
from sys import argv
from sys import exit

from modules.llvm import LLVMBuildConfig
from modules.llvm import LLVMSubproject
from modules.llvm import LLVMSourceConfig
from modules.llvm import run_test_suite
from modules.llvm import setup_test_suite
from modules.utils import CommandPrinter
from modules.utils import CommandRunner
from modules.utils import get_remote_branch
from modules.utils import push_branch

from linaropy.cd import cd
from linaropy.git.clone import Clone
from linaropy.proj import Proj

from argparse import Action, ArgumentParser, RawTextHelpFormatter
from functools import partial


def die(message, config_to_dump=None):
    """Print an error message and exit."""
    print(message)

    if config_to_dump is not None:
        dump_config(config_to_dump)

    exit(1)


def dump_config(config):
    """Dump the list of projects that are enabled in the given config."""

    print("Projects linked:")
    enabled = config.get_enabled_subprojects()
    if not enabled:
        print("none")
    else:
        for subproj in sorted(enabled):
            print("  + {}".format(subproj))


def subproj_to_repo_map(subprojs, proj, reposRoot, dry=False):
    """Get a dictionary mapping each subproject in subprojs to its repo."""
    subprojsToRepos = {}
    repo = None
    for subproj in subprojs:
        if not dry:
            repo = Clone(proj, os.path.join(reposRoot, subproj))
        subprojsToRepos[subproj] = repo

    return subprojsToRepos


def projects(args):
    """Add/remove subprojects based on the values in args."""

    proj = Proj()

    llvm_worktree_root = args.sources
    llvm_repos_root = args.repos
    config = LLVMSourceConfig(proj, llvm_worktree_root, dry=False)

    if not args.add and not args.remove:
        # Nothing to change, just print the current configuration
        dump_config(config)
        exit(0)

    to_add = {}
    if args.add:
        to_add = subproj_to_repo_map(args.add, proj, llvm_repos_root)

    try:
        config.update(to_add, args.remove)
    except (EnvironmentError, ValueError) as exc:
        die("Failed to update subprojects because:\n{}".format(str(exc)))

    dump_config(config)


def push_current_branch(args):
    """Push current branch to origin."""

    proj = Proj()

    llvm_worktree_root = args.sources
    config = LLVMSourceConfig(proj, llvm_worktree_root, dry=False)

    llvm_worktree = Clone(proj, llvm_worktree_root)
    local_branch = llvm_worktree.getbranch()

    try:
        remote_branch = get_remote_branch(llvm_worktree, local_branch)
        config.for_each_enabled(partial(push_branch, proj, local_branch,
                                        remote_branch))
        print("Pushed to {}".format(remote_branch))
    except (EnvironmentError, RuntimeError) as exc:
        die("Failed to push branch because: " + str(exc) + str(exc.__cause__))


def cmake_flags_from_args(defs):
    """
    Get a list of valid CMake flags from the input VAR=VALUE list.
    This boils down to just adding -D in front of each of them.
    """
    return ["-D{}".format(v) for v in defs]


def configure_build(args):
    """Configure a given build directory."""

    proj = Proj()

    llvm_worktree_root = args.sources
    sourceConfig = LLVMSourceConfig(proj, llvm_worktree_root, args.dry)

    if args.dry:
        consumer = CommandPrinter()
    else:
        if not os.path.exists(args.build):
            os.makedirs(args.build)
        consumer = CommandRunner()

    buildConfig = LLVMBuildConfig(sourceConfig, args.build, consumer)

    if args.defs:
        args.defs = cmake_flags_from_args(args.defs)

    try:
        buildConfig.cmake(args.defs, args.generator)
    except RuntimeError as exc:
        die("Failed to configure {} because:\n{}".format(args.build, str(exc)))


def run_build(args):
    """Run a build command in a given directory."""
    build_dir = args.build

    if args.dry:
        consumer = CommandPrinter()
    else:
        consumer = CommandRunner()

    try:
        LLVMBuildConfig(None, args.build, consumer).build(args.flags)
    except RuntimeError as exc:
        die("Failed to build {} because:\n{}".format(args.build, str(exc)))


def setup_the_test_suite(args):
    """Setup a sandbox for the test-suite."""
    if args.dry:
        consumer = CommandPrinter()
    else:
        consumer = CommandRunner()

    try:
        setup_test_suite(consumer, args.sandbox, args.lnt)
    except RuntimeError as exc:
        die("Failed to setup the test-suite because:\n{}".format(str(exc)))


def run_the_test_suite(args):
    """Run the test-suite in a given sandbox."""
    if args.dry:
        consumer = CommandPrinter()
    else:
        consumer = CommandRunner()

    compilers = ["--cc={}".format(args.cc)]
    if args.cxx:
        compilers.append("--cxx={}".format(args.cxx))

    try:
        run_test_suite(consumer, args.sandbox, args.testsuite, args.lit,
                       compilers + args.flags)
    except RuntimeError as exc:
        die("Failed to run the test-suite because:\n{}".format(str(exc)))


def build_and_test(args):
    """
    Build a set of LLVM subprojects, in one or two stages, with or without a
    test-suite run.
    """

    proj = Proj()

    dryRun = args.dry
    llvmRepos = args.repos
    llvmWorktreeRoot = args.sources

    stage1BuildDir = args.stage1
    stage1Subprojs = args.stage1Subprojs
    stage1Defs = cmake_flags_from_args(args.stage1Defs)
    stage1BuildFlags = args.stage1BuildFlags

    stage2BuildDir = args.stage2
    stage2Subprojs = args.stage2Subprojs
    stage2Defs = cmake_flags_from_args(args.stage2Defs)
    stage2BuildFlags = args.stage2BuildFlags

    enableTestSuite = args.enableTestSuite
    sandboxDir = args.sandbox
    testSuiteFlags = args.testSuiteFlags

    if not "clang" in stage1Subprojs:
        if enableTestSuite:
            die("Can't enable the test-suite if stage 1 doesn't build clang."
                "Please add clang to the subprojects to be built for stage 1.")

        if stage2BuildDir:
            die("Can't enable stage 2 if stage 1 doesn't build clang."
                "Please add clang to the subprojects to be built for stage 1.")

    if dryRun:
        consumer = CommandPrinter()
    else:
        consumer = CommandRunner()

    try:
        sourceConfig = LLVMSourceConfig(proj, llvmWorktreeRoot, args.dry)
        if stage1Subprojs:
            # FIXME: Decide whether or not we want to remove anything that isn't
            # in stage1Subprojs (in case there are already some enabled
            # subprojects in the source config).
            sourceConfig.update(
                subproj_to_repo_map(stage1Subprojs, proj, llvmRepos,
                                    args.dry),
                [])

        if not dryRun and not os.path.exists(stage1BuildDir):
            os.makedirs(stage1BuildDir)

        buildConfig1 = LLVMBuildConfig(sourceConfig, stage1BuildDir, consumer)
        buildConfig1.cmake(stage1Defs, "Ninja")
        buildConfig1.build(stage1BuildFlags)
        testedBuildDir = stage1BuildDir

        if stage2BuildDir is not None:
            if stage2Subprojs:
                toAdd = list(set(stage2Subprojs) - set(stage1Subprojs))
                toRemove = list(set(stage1Subprojs) - set(stage2Subprojs))
                sourceConfig.update(
                    subproj_to_repo_map(toAdd, proj, llvmRepos, args.dry),
                    toRemove)

            if not dryRun and not os.path.exists(stage2BuildDir):
                os.makedirs(stage2BuildDir)

            buildConfig2 = LLVMBuildConfig(sourceConfig, stage2BuildDir,
                                           consumer)

            buildConfig2.cmake(
                stage2Defs + [
                    "-DCMAKE_C_COMPILER={}/bin/clang".format(stage1BuildDir),
                    "-DCMAKE_CXX_COMPILER={}/bin/clang++".format(stage1BuildDir)],
                "Ninja")
            buildConfig2.build(stage2BuildFlags)
            testedBuildDir = stage2BuildDir

        if enableTestSuite:
            testSuiteDir = os.path.join(llvmRepos, "test-suite")
            lntDir = os.path.join(llvmRepos, "lnt")

            setup_test_suite(consumer, sandboxDir, lntDir)

            lit = os.path.join(testedBuildDir, "bin", "llvm-lit")
            clang = os.path.join(testedBuildDir, "bin", "clang")
            run_test_suite(consumer, sandboxDir, testSuiteDir, lit,
                           ["--cc={}".format(clang)] + testSuiteFlags)

    except RuntimeError as exc:
        die("Failed because:\n{}".format(str(exc)))


##########################################################################
# Command line parsing                                                   #
##########################################################################

# If we decide we want shorthands for the subprojects, we can append to this
# list
valid_subprojects = list(LLVMSubproject.get_all_subprojects().keys())

options = ArgumentParser(formatter_class=RawTextHelpFormatter)
subcommands = options.add_subparsers(dest="subcommand")

# Subcommand for adding / removing subprojects
projs = subcommands.add_parser(
    "projects", help="Add/remove LLVM subprojects.\n"
                     "Adding a subproject will create a worktree for it "
                     "somewhere in the LLVM source tree, on the same git "
                     "branch as LLVM itself.\n"
                     "Removing a subproject will remove the worktree, but "
                     "not the underlying git branch.")
projs.set_defaults(run_command=projects)

# TODO: Overwriting previous values is not necessarily what users expect (so for
# instance --add S1 S2 --remove S3 --add S4 would lead to adding only S4). We
# can do better by using action='append', which would create a list (of lists?
# or of lists and scalars?) that we can flatten to obtain all the values passed
# by the user.
projs.add_argument(
    '-a', '--add',
    nargs='+',
    choices=valid_subprojects,
    metavar='subproject',
    help="Enable given subprojects. Valid values are:\n\t{}\n".format(
         "\n\t".join(valid_subprojects)))
projs.add_argument(
    '-r', '--remove',
    nargs='+',
    choices=valid_subprojects,
    metavar='subproject',
    help="Disable given subprojects.")
projs.add_argument(
    '--repos',
    help="Path to the directory containing the repositories for all LLVM "
         "subprojects.")
projs.add_argument(
    '--source-dir',
    dest='sources',
    required=True,
    help="Path to the directory containing the LLVM worktree that we're adding "
         "or removing subprojects from.")

# Subcommand for pushing the current branch to origin
push = subcommands.add_parser(
    "push",
    help="Push current branch to origin linaro-local/<user>/<branch>, "
         "for all enabled subprojects.")
push.set_defaults(run_command=push_current_branch)
push.add_argument(
    '--source-dir',
    dest='sources',
    required=True,
    help="Path to the directory containing the LLVM worktree.")

# Subcommand for configuring a build directory
configure = subcommands.add_parser(
    'configure',
    help="Run CMake in the given build directory.")
configure.add_argument(
    '--source-dir',
    dest='sources',
    required=True,
    help="Path to the sources directory. It should contain an LLVM worktree.")
configure.add_argument(
    '--build-dir',
    dest='build',
    required=True,
    help="Path to the build directory. It will be created if it does not exist")
configure.add_argument(
    '--cmake-generator',
    dest='generator',
    default='Ninja',
    help="CMake generator to use (default is Ninja).")
configure.add_argument(
    '--cmake-def',
    dest='defs',
    metavar='VAR=VALUE',
    default=[],
    action='append',
    # We add the -D in front of the variable ourselves because the argument
    # parsing gets confused otherwise (and quoting doesn't help).
    help="Additional CMake definitions, e.g. CMAKE_BUILD_TYPE=Release."
    "May be passed several times. The -D is added automatically.")
configure.add_argument(
    '-n', '--dry-run',
    dest='dry',
    action='store_true',
    default=False,
    help="Print the CMake command instead of executing it.")
configure.set_defaults(run_command=configure_build)

# Subcommand for building a target
build = subcommands.add_parser(
    'build',
    help="Run a build command in the given directory."
    "The build command can be either a 'ninja' or a 'make' command, depending "
    "on what the build directory contains. First, we look for a 'build.ninja' "
    "file. If that is not found, we look for a 'Makefile'. If that is not "
    "found either, the script fails.")
build.add_argument(
    '--build-dir',
    dest='build',
    required=True,
    help="Path to the build directory. It must have already been configured.")
build.add_argument(
    '-n', '--dry-run',
    dest='dry',
    action='store_true',
    default=False,
    help="Print the build command instead of executing it.")
build.add_argument(
    '--build-flag',
    dest='flags',
    metavar='FLAG',
    default=[],
    action='append',
    help="Additional flags for the build command (e.g. targets to build). "
    "May be passed several times. If your flag starts with a '-', use "
    "'--build-flag=-FLAG' to pass it.")
build.set_defaults(run_command=run_build)

# Subcommand for setting up the test-suite
setupTestSuite = subcommands.add_parser(
    'setup-test-suite',
    help="Prepare a sandbox for running the test-suite.")
setupTestSuite.add_argument(
    '--sandbox',
    required=True,
    help="Path where we should setup the sandbox.")
setupTestSuite.add_argument(
    '--lnt',
    required=True,
    help="Path to the LNT sources.")
setupTestSuite.add_argument(
    '-n', '--dry-run',
    dest='dry',
    action='store_true',
    default=False,
    help="Print the commands instead of executing them.")
setupTestSuite.set_defaults(run_command=setup_the_test_suite)

# Subcommand for running the test-suite
runTestSuite = subcommands.add_parser(
    'run-test-suite',
    help="Run the test-suite in the given sandbox.")
runTestSuite.add_argument(
    '--sandbox',
    required=True,
    help="Path to the sandbox. It must point to a virtualenv with a LNT setup.")
runTestSuite.add_argument(
    '--test-suite',
    dest="testsuite",
    required=True,
    help="Path to the test-suite repo.")
runTestSuite.add_argument(
    '--use-lit',
    dest="lit",
    required=True,
    help="Path to llvm-lit.")
runTestSuite.add_argument(
    '--lnt-flag',
    dest='flags',
    metavar='FLAG',
    default=[],
    action='append',
    help="Additional flags to be passed to LNT when running the test-suite."
    "May be passed several times. If your flag starts with a '-', use "
    "'--lnt-flag=-FLAG' to pass it.")
runTestSuite.add_argument(
    # We can pass --cc through the --lnt-flag interface, but we generally won't
    # want to test the system compiler, so force the user to be specific.
    '--cc',
    required=True,
    help="The path to the C compiler that we're testing.")
runTestSuite.add_argument(
    # For symmetry, we also provide a --cxx argument, but this one isn't
    # required since LNT tries to guess it based on the value of --cc.
    '--cxx',
    required=False,
    help="The path to the C++ compiler that we're testing.")
runTestSuite.add_argument(
    '-n', '--dry-run',
    dest='dry',
    action='store_true',
    default=False,
    help="Print the commands instead of executing them.")
runTestSuite.set_defaults(run_command=run_the_test_suite)

buildAndTest = subcommands.add_parser(
    'build-and-test',  # TODO: This really needs a better name...
    fromfile_prefix_chars='@',
    help="Run complex build scenarios with one or two stages of clang and "
         "optionally a test-suite run. This should be flexible enough to allow "
         "us to reproduce any buildbot configuration, but it can obviously be "
         "used for other purposes as well.")
buildAndTest.set_defaults(run_command=build_and_test)
buildAndTest.add_argument(
    '--repos-dir',
    dest='repos',
    required=True,
    help="Path to the root directory containing the repositories for LLVM and "
         "the other subprojects.")
buildAndTest.add_argument(
    '--source-dir',
    dest='sources',
    required=True,
    help="Path to the directory containing the LLVM worktree that we're going "
         "to build from.")
buildAndTest.add_argument(
    '--stage1-build-dir',
    dest='stage1',
    required=True,
    help="Path to the build directory for stage 1.")
buildAndTest.add_argument(
    '--stage1-subproject',
    dest='stage1Subprojs',
    metavar='SUBPROJ',
    choices=valid_subprojects,
    default=[],
    action='append',
    help="Subprojects to enable for stage 1 of the build. Can be passed "
         "multiple times. Valid values for the subproject are: {}. "
         "If this is a 2-stage build, the same subprojects will be used for "
         "both stages unless other subprojects are explicitly requested for "
         "stage 2.".format(" ".join(valid_subprojects)))
buildAndTest.add_argument(
    '--stage1-cmake-def',
    dest='stage1Defs',
    metavar='VAR=VALUE',
    default=[],
    action='append',
    help="Additional CMake definitions for stage 1, e.g. "
         "CMAKE_BUILD_TYPE=Release. Can be passed multiple times. "
         "The -D is added automatically. Does not affect stage 2.")
buildAndTest.add_argument(
    '--stage1-build-flag',
    dest='stage1BuildFlags',
    metavar='FLAG',
    default=[],
    action='append',
    help="Additional flags for the stage 1 build command (e.g. targets to "
         "build). Can be passed multiple times. If your flag starts with "
         "a '-', use '--stage1-build-flag=-FLAG' to pass it. "
         "Does not affect stage 2.")
buildAndTest.add_argument(
    '--stage2-build-dir',
    dest='stage2',
    help="Path to the build directory for stage 2.")
buildAndTest.add_argument(
    '--stage2-subproject',
    dest='stage2Subprojs',
    metavar='SUBPROJ',
    choices=valid_subprojects,
    default=[],
    action='append',
    help="Subprojects to enable for stage 2 of the build. Can be passed "
         "multiple times. Valid values for the subproject are: {}. "
         "If this is a 2-stage build, the same subprojects will be used for "
         "both stages unless other subprojects are explicitly requested for "
         "stage 2.".format(" ".join(valid_subprojects)))
buildAndTest.add_argument(
    '--stage2-cmake-def',
    dest='stage2Defs',
    metavar='VAR=VALUE',
    default=[],
    action='append',
    help="Additional CMake definitions for stage 2, e.g. "
         "CMAKE_BUILD_TYPE=Release. Can be passed multiple times. "
         "The -D is added automatically.")
buildAndTest.add_argument(
    '--stage2-build-flag',
    dest='stage2BuildFlags',
    metavar='FLAG',
    default=[],
    action='append',
    help="Additional flags for the stage 2 build command (e.g. targets to "
         "build). Can be passed multiple times. If your flag starts with "
         "a '-', use '--stage2-build-flag=-FLAG' to pass it.")
buildAndTest.add_argument(
    "--enable-test-suite",
    dest='enableTestSuite',
    action='store_true',
    default=False,
    help="Whether or not to run the test-suite with the last compiler built.")
buildAndTest.add_argument(
    "--sandbox",
    help="Path to the sandbox where the test-suite should be setup.")
buildAndTest.add_argument(
    '--lnt-flag',
    dest='testSuiteFlags',
    metavar='FLAG',
    default=[],
    action='append',
    help="Additional flags to be passed to LNT when running the test-suite."
    "May be passed several times. If your flag starts with a '-', use "
    "'--lnt-flag=-FLAG' to pass it.")
buildAndTest.add_argument(
    '-n', '--dry-run',
    dest='dry',
    action='store_true',
    default=False,
    help="Print the commands instead of executing them.")

args = options.parse_args()
if args.subcommand == "projects" and args.add and not args.repos:
    projs.error(
        "When adding a subproject you must also pass the --repos argument")
args.run_command(args)