# Copyright 2021-2022 Timo Röhling <roehling@debian.org>
# SPDX-License-Identifier: MIT
import argparse
import logging
import os
from pathlib import Path
import re
import subprocess
import shlex
import shutil
import sys
from typing import cast, Optional, Dict, List, Set

from ..common import (
    Package,
    PyVerTuple,
    find_packages,
    compute_python_package_path,
    compute_ament_cmake_python_path,
    sort_packages,
    get_build_var,
    get_python_versions,
    is_default_python,
    query_dpkg_architecture,
)

# command line identation
cmdind = "    "
destdirs: Set[Path] = set()
catkin_srcs: List[str] = []
dpkg_info = query_dpkg_architecture()


def prepare_argparse() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser()
    p.add_argument(
        "--dry-run", action="store_true", help="do not execute any commands for real"
    )
    p.add_argument(
        "--sourcedir",
        default=Path("."),
        type=Path,
        help="root directory of the source package",
    )
    g = p.add_mutually_exclusive_group()
    g.add_argument(
        "--search-depth",
        default=2,
        type=int,
        help="limit maximum depth when searching for packages recursively",
    )
    g.add_argument(
        "--unlimited-search-depth",
        action="store_const",
        dest="search_depth",
        const=-1,
        help="do not limit maximum depth when searching for packages",
    )
    p.add_argument(
        "--builddir", default=Path(".rosbuild"), type=Path, help="build directory"
    )
    p.add_argument(
        "--destdir",
        default=Path("debian/tmp"),
        type=Path,
        help="install staging directory",
    )
    p.add_argument(
        "--install-prefix", default=Path("/usr"), type=Path, help="install prefix"
    )
    p.add_argument("--verbose", action="store_true", help="make verbose output")
    g = p.add_mutually_exclusive_group()
    g.add_argument(
        "--detect", action="store_true", help="detect if ROS packages are to be built"
    )
    g.add_argument(
        "--build-types", action="store_true", help="list detected build types"
    )
    g.add_argument(
        "--build-order", action="store_true", help="list packages in build order"
    )
    g.add_argument("--clean", action="store_true", help="clean source tree and quit")
    return p


def main() -> int:
    p = prepare_argparse()
    args = p.parse_args()
    return run(args, cast(Dict[str, str], os.environ))


def run(args: argparse.Namespace, environ: Dict[str, str]) -> int:
    # initialize script
    logging.basicConfig(
        format="%(levelname).1s: dh_ros %(module)s:%(lineno)d: %(message)s"
    )
    log = logging.getLogger("dhros")
    if environ.get("DH_VERBOSE") == "1":
        args.verbose = True
    if args.verbose:
        log.setLevel(logging.DEBUG)
    else:
        log.setLevel(logging.INFO)

    packages = find_packages(args.sourcedir, depth=args.search_depth)

    if args.detect:
        return 0 if packages else 1

    if args.build_types:
        sys.stdout.write(
            "\n".join(sorted(set(cast(str, p.get_build_type()) for _, p in packages)))
        )
        return 0

    buildopts: List[str] = re.split(
        r"[ ,\t\r\n]+", environ.get("DEB_BUILD_OPTIONS", "")
    )

    packages = sort_packages(packages, test_depends="nocheck" not in buildopts)
    skipped_packages = get_build_var(None, "ROS_SKIP_PACKAGES", env=environ)
    skipped_packages = set(shlex.split(skipped_packages) if skipped_packages else [])
    skipped_tests = get_build_var(None, "ROS_SKIP_TESTS", env=environ)
    if skipped_tests == "1":
        skipped_tests = set(p[1].name for p in packages)
    else:
        skipped_tests = set(shlex.split(skipped_tests) if skipped_tests else [])
    python_extension_packages = get_build_var(
        None, "ROS_PYTHON_EXTENSION_PACKAGES", env=environ
    )
    if python_extension_packages == "1":
        python_extension_packages = set(p[1].name for p in packages)
    else:
        python_extension_packages = set(
            shlex.split(python_extension_packages) if python_extension_packages else []
        )
    ignore_test_results = get_build_var(None, "ROS_IGNORE_TEST_RESULTS", env=environ)
    if ignore_test_results == "1":
        ignore_test_results = set(p[1].name for p in packages)
    else:
        ignore_test_results = set(
            shlex.split(ignore_test_results) if ignore_test_results else []
        )
    install_prefix = Path(
        get_build_var(None, "ROS_INSTALL_PREFIX", str(args.install_prefix), env=environ)
    ).resolve()

    if args.build_order:
        sys.stdout.write("\n".join(p.name for _, p in packages))
        return 0

    if args.clean:
        try:
            for path, package in packages:
                if package.name not in skipped_packages:
                    do_action(
                        "clean",
                        args,
                        path,
                        package,
                        buildopts,
                        install_prefix,
                        environ,
                        None,
                    )
            if args.builddir.is_dir():
                log.info(f"Removing {args.builddir}")
                shutil.rmtree(args.builddir, ignore_errors=True)
        except KeyError as e:
            log.error(str(e))
            return 1
        except subprocess.CalledProcessError as e:
            return e.returncode or 1
        return 0

    sys.stdout.write(
        "#############################################################################\n"
        "## Detected ROS packages (in build order):                                 ##\n"
        "##                                                                         ##\n"
    )
    for path, package in packages:
        line = f"- {package.name} [{package.get_build_type()}]"
        sys.stdout.write(f"## {line:<71} ##\n")
    sys.stdout.write(
        "#############################################################################\n"
    )

    try:
        global destdirs
        destdirs = set()
        for path, package in packages:
            if package.name in skipped_packages:
                log.info(f"Skipping ROS package {package.name}")
                continue

            sys.stdout.write(
                "\n"
                "=============================================================================\n"
                f"= ROS Package {package.name:<61} =\n"
                "=============================================================================\n"
            )
            if package.get_build_type() == "catkin":
                catkin_srcs.append(str(path.resolve()))
            python_versions = [None]
            package_buildopts: List[str] = []
            if package.name in skipped_tests:
                package_buildopts.append("nocheck")
            if (
                package.name in python_extension_packages
                and package.get_build_type()
                in [
                    "cmake",
                    "catkin",
                    "ament_cmake",
                ]
            ):
                default_python, supported_python = get_python_versions()
                if default_python is not None and supported_python is not None:
                    python_versions = [
                        tuple(v) for v in supported_python if v != default_python
                    ]
                    python_versions.append(tuple(default_python))
                    s = [".".join(str(v) for v in pv) for pv in python_versions]
                    log.info(
                        f"Python extensions will be built against Python {', '.join(s)}"
                    )

            for python_version in python_versions:
                do_action(
                    "configure",
                    args,
                    path,
                    package,
                    buildopts + package_buildopts,
                    install_prefix,
                    environ,
                    python_version,
                )
                do_action(
                    "build",
                    args,
                    path,
                    package,
                    buildopts,
                    install_prefix,
                    environ,
                    python_version,
                )
                if "nocheck" not in buildopts and package.name not in skipped_tests:
                    try:
                        do_action(
                            "test",
                            args,
                            path,
                            package,
                            buildopts,
                            install_prefix,
                            environ,
                            python_version,
                        )
                    except subprocess.CalledProcessError as e:
                        if package.name in ignore_test_results:
                            log.info(
                                f"Ignoring test failure in ROS package {package.name}"
                            )
                        else:
                            raise
                else:
                    log.info(f"Skipping tests for ROS package {package.name}")
                do_action(
                    "install",
                    args,
                    path,
                    package,
                    buildopts,
                    install_prefix,
                    environ,
                    python_version,
                )
        sys.stdout.write(
            "\n"
            "#############################################################################\n"
            "## All ROS packages have been built successfully                           ##\n"
            "#############################################################################\n"
        )
    except KeyError as e:
        log.error(str(e))
        return 1
    except subprocess.CalledProcessError as e:
        return e.returncode or 1
    return 0


def do_action(
    action: str,
    args: argparse.Namespace,
    path: Path,
    package: Package,
    buildopts: List[str],
    install_prefix: Path,
    environ: Dict[str, str],
    python_version: Optional[PyVerTuple],
) -> None:
    global destdirs
    global catkin_srcs
    log = logging.getLogger("dhros")
    destdir = Path(
        get_build_var(package.name, "ROS_DESTDIR", str(args.destdir), env=environ)
    ).resolve()
    homedir = Path(
        get_build_var(
            package.name, "ROS_HOME", str(args.builddir / ".roshome"), env=environ
        )
    ).resolve()
    builddir = args.builddir / package.name
    python_interpreter = ""
    python_version_str = ""
    if python_version is not None:
        python_version_str = ".".join(str(v) for v in python_version)
        builddir = builddir / f"py{python_version_str}"
        if is_default_python(python_version):
            python_interpreter = f"/usr/bin/python{python_version[0]}"
        else:
            python_interpreter = f"/usr/bin/python{python_version_str}"
    catkin_test_results_dir = Path(
        get_build_var(
            package.name,
            "ROS_CATKIN_TEST_RESULTS_DIR",
            str(builddir / "test_results"),
            env=environ,
        )
    ).resolve()
    catkin_devel_dir = (args.builddir / "_devel").resolve()
    relative_install_prefix = install_prefix.relative_to("/")
    destdirs.add(destdir)
    package_suffix = re.sub(r"[^A-Za-z0-9]", "_", package.name)
    build_type = cast(str, package.get_build_type())
    hook_vars = {
        "build_dir": shlex.quote(str(builddir.resolve())),
        "build_type": shlex.quote(build_type),
        "destdir": shlex.quote(str(destdir)),
        "dir": shlex.quote(str(path.resolve())),
        "homedir": shlex.quote(str(homedir)),
        "package": shlex.quote(package.name),
        "package!id": package_suffix,
        "prefix": shlex.quote(str(install_prefix)),
        "version": shlex.quote(package.version),
        "python_version": python_version_str,
        "python_interpreter": python_interpreter,
    }
    env = environ.copy()
    pybuild_per_package_keys = [
        key
        for key in env
        if key.startswith("PYBUILD_") and key.endswith(f"_{package_suffix}")
    ]
    for key in pybuild_per_package_keys:
        env[key[: -len(package_suffix) - 1]] = env[key]
        del env[key]
    env["ROS_HOME"] = str(homedir)
    if install_prefix == Path("/usr") and "DEB_PYTHON_INSTALL_LAYOUT" not in env:
        env["DEB_PYTHON_INSTALL_LAYOUT"] = "deb"
    if install_prefix != Path("/usr"):
        pybuild_install_args = f"--prefix={shlex.quote(str(install_prefix))} --install-lib=$base/lib/{{interpreter}}/site-packages --install-scripts=$base/bin"
        pybuild_user_install_args = env.get("PYBUILD_INSTALL_ARGS", "")
        env["PYBUILD_INSTALL_ARGS"] = (
            pybuild_install_args + " " + pybuild_user_install_args
        )

    if action == "configure" and build_type == "ament_python":
        if not args.dry_run:
            pybuild_dir = args.sourcedir / ".pybuild"
            if pybuild_dir.is_dir():
                log.info(f"Removing {pybuild_dir}")
                shutil.rmtree(pybuild_dir, ignore_errors=True)

    prefixes = env.get("CMAKE_PREFIX_PATH", None)
    prefixes = prefixes.split(":") if prefixes is not None else []
    if install_prefix != Path("/usr"):
        prefixes.append(str(install_prefix))
    env["CMAKE_PREFIX_PATH"] = ":".join(
        prefixes + [str(d / relative_install_prefix) for d in destdirs]
    )
    if build_type == "ament_cmake":
        prefixes = env.get("AMENT_PREFIX_PATH", None)
        prefixes = prefixes.split(":") if prefixes is not None else []
        if install_prefix != Path("/usr"):
            prefixes.append(str(install_prefix))
        env["AMENT_PREFIX_PATH"] = ":".join(
            prefixes + [str(d / relative_install_prefix) for d in destdirs]
        )
    if build_type == "catkin":
        package_paths = env.get("ROS_PACKAGE_PATH", None)
        package_paths = package_paths.split(":") if package_paths is not None else []
        if install_prefix != Path("/usr"):
            package_paths.append(str(install_prefix / "share"))
        env["ROS_PACKAGE_PATH"] = ":".join(
            package_paths
            + catkin_srcs
            + [str(d / relative_install_prefix / "share") for d in destdirs]
        )

    pythonpaths = env.get("PYTHONPATH", None)
    pythonpaths = pythonpaths.split(":") if pythonpaths is not None else []
    if build_type == "ament_cmake":
        pythonpaths += compute_ament_cmake_python_path(builddir)
    if build_type == "catkin":
        pythonpaths += compute_python_package_path(catkin_devel_dir)
    for d in destdirs:
        pythonpaths += compute_python_package_path(d / relative_install_prefix)
    if install_prefix != Path("/usr"):
        pythonpaths += compute_python_package_path(install_prefix)
    env["PYTHONPATH"] = ":".join(pythonpaths)

    deb_host_multiarch = dpkg_info.get("DEB_HOST_MULTIARCH", "")
    deb_build_multiarch = dpkg_info.get("DEB_BUILD_MULTIARCH", "")
    if (
        "_PYTHON_SYSCONFIGDATA_NAME" not in env
        and deb_host_multiarch != deb_build_multiarch
    ):
        deb_host_arch_os = dpkg_info.get("DEB_HOST_ARCH_OS", "")
        env["_PYTHON_SYSCONFIGDATA_NAME"] = (
            f"_sysconfigdata__{deb_host_arch_os}_{deb_host_multiarch}"
        )

    if build_type == "catkin":
        cpath = env.get("CPATH", None)
        cpath = cpath.split(":") if cpath is not None else []
        for d in destdirs:
            cpath.append(str(d / "usr" / "include"))
            if deb_host_multiarch:
                cpath.append(str(d / "usr" / "include" / deb_host_multiarch))
        env["CPATH"] = ":".join(cpath)

    if action == "test":
        if "ROS_IP" not in env and "ROS_HOSTNAME" not in env:
            env["ROS_IP"] = "127.0.0.1"

    if args.verbose:
        for key, value in env.items():
            log.debug(f"env {key}={value!r}")

    execute_hook(f"before_{action}", package, hook_vars, env, dry_run=args.dry_run)
    execute_hook(
        f"before_{action}_{build_type}", package, hook_vars, env, dry_run=args.dry_run
    )

    if not execute_hook(
        f"custom_{action}_{build_type}", package, hook_vars, env, dry_run=args.dry_run
    ) and not execute_hook(
        f"custom_{action}", package, hook_vars, env, dry_run=args.dry_run
    ):
        cmdline = [f"dh_auto_{action}", f"--sourcedir={path}"]
        if build_type in ["ament_python"]:
            cmdline += ["--buildsystem=pybuild"]
        elif build_type in ["ament_cmake", "catkin", "cmake"]:
            cmdline += [
                f"--builddir={builddir}",
                "--buildsystem=cmake",
            ]
        if action == "install":
            cmdline += [f"--destdir={str(destdir)}"]

        cmdline += ["--"]

        if action == "configure":
            if build_type in ["ament_cmake", "catkin", "cmake"]:
                build_testing = "OFF" if "nocheck" in buildopts else "ON"
                cmdline += [
                    "--no-warn-unused-cli",
                    "-DBUILD_SHARED_LIBS=ON",
                    f"-DBUILD_TESTING={build_testing}",
                    "-DCMAKE_POLICY_VERSION_MINIMUM=3.10",
                ]
                if install_prefix != Path("/usr"):
                    cmdline.append(f"-DCMAKE_INSTALL_PREFIX={install_prefix}")
            if build_type == "catkin":
                cmdline += [
                    "-DCATKIN_BUILD_BINARY_PACKAGE=ON",
                    f"-DCATKIN_DEVEL_PREFIX={str(catkin_devel_dir)}",
                    f"-DCATKIN_TEST_RESULTS_DIR={str(catkin_test_results_dir)}",
                ]
            if build_type == "ament_cmake":
                cmdline += ["-DAMENT_LINT_AUTO=OFF"]
                if install_prefix == Path("/usr"):
                    cmdline += [
                        "-DAMENT_CMAKE_ENVIRONMENT_PACKAGE_GENERATION=OFF",
                        "-DAMENT_CMAKE_ENVIRONMENT_PARENT_PREFIX_PATH_GENERATION=OFF",
                    ]
            if python_version is not None:
                cmdline += [
                    f"-DPYTHON_EXECUTABLE={python_interpreter}",
                    f"-DPython_EXECUTABLE={python_interpreter}",
                    f"-DPython{python_version[0]}_EXECUTABLE={python_interpreter}",
                ]
        elif action == "test" and build_type == "catkin":
            cmdline = [
                "cmake",
                "--build",
                str(builddir),
                "--target",
                "run_tests",
            ]

        for extra_args_var in [
            f"ROS_{action.upper()}_{build_type.upper()}_WRAPPER",
            f"ROS_{action.upper()}_WRAPPER",
        ]:
            extra_args = get_build_var(package.name, extra_args_var, env=environ)
            if extra_args:
                try:
                    cmdline = shlex.split(extra_args.format(**hook_vars)) + cmdline
                except KeyError as e:
                    log.error(
                        f"Invalid substitution {{{str(e)}}} in {extra_args_var}: {extra_args!r}"
                    )
                    raise
                break
        for extra_args_var in [
            f"ROS_{action.upper()}_ARGS",
            f"ROS_{action.upper()}_{build_type.upper()}_ARGS",
        ]:
            extra_args = get_build_var(package.name, extra_args_var, env=environ)
            if extra_args:
                try:
                    cmdline += shlex.split(extra_args.format(**hook_vars))
                except KeyError as e:
                    log.error(
                        f"Invalid substitution {{{str(e)}}} in {extra_args_var}: {extra_args!r}"
                    )
                    raise

        sys.stdout.write(cmdind)
        sys.stdout.write(shlex.join(cmdline))
        sys.stdout.write("\n")
        sys.stdout.flush()
        if not args.dry_run:
            subprocess.check_call(cmdline, env=env)

            if (
                action == "test"
                and build_type == "catkin"
                and package.name != "catkin"
                and catkin_test_results_dir.is_dir()
            ):
                cmdline = ["catkin_test_results"]
                if args.verbose:
                    cmdline.append("--verbose")
                cmdline.append(str(catkin_test_results_dir))
                sys.stdout.write(f"{cmdind}{shlex.join(cmdline)}\n")
                sys.stdout.flush()
                subprocess.check_call(cmdline, env=env)

    execute_hook(
        f"after_{action}_{build_type}", package, hook_vars, env, dry_run=args.dry_run
    )
    execute_hook(f"after_{action}", package, hook_vars, env, dry_run=args.dry_run)
    if not args.dry_run:
        if action == "install":
            for d in compute_python_package_path(destdir / relative_install_prefix):
                d = Path(d) / ".pytest_cache"
                if d.is_dir():
                    log.info(f"Removing {d}")
                    shutil.rmtree(d, ignore_errors=True)
        if action == "clean":
            for pycache in path.glob("**/__pycache__/"):
                log.info(f"Removing {pycache}")
                shutil.rmtree(pycache, ignore_errors=True)


def execute_hook(
    hook: str,
    package: Package,
    vars: Dict[str, str],
    env: Dict[str, str],
    dry_run: bool = False,
):
    log = logging.getLogger("dhros")
    cmdline = get_build_var(package.name, f"ROS_EXECUTE_{hook.upper()}", env=env)
    if cmdline is None:
        log.debug(f"Nothing to be done for hook {hook!r}")
        return False

    try:
        cmdline = cmdline.format(**vars)
    except KeyError as e:
        log.error(f"Invalid substitution {{{str(e)}}} in hook {hook!r}")
        raise

    sys.stdout.write(cmdind)
    sys.stdout.write(cmdline)
    sys.stdout.write("\n")
    sys.stdout.flush()
    if not dry_run:
        subprocess.check_call(["/bin/sh", "-c", cmdline], env=env)
    return True
