#!/usr/bin/env python3

# Copyright Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT

from __future__ import print_function

import argparse
import grp
import os
import subprocess
from pathlib import Path

dockerfiles = {
    "clang": "docker/dockerfile-ubuntu-clang",
    "gcc": "docker/dockerfile-ubuntu-gcc",
}


def gitTop() -> Path:
    command = ["git"] + ["rev-parse", "--show-toplevel"]
    p = subprocess.run(
        command,
        check=True,
        capture_output=True,
        text=True,
    )
    return Path(p.stdout.strip()).resolve()


def parseArgs():
    parser = argparse.ArgumentParser(
        description="Build and Run Project Docker Containers. "
        "Run from root of git directory.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )

    parser.add_argument(
        "--user", default=os.environ["USER"], help="Username to create in containers."
    )
    parser.add_argument(
        "--uid", default=str(os.getuid()), help="UID to assign to the created user."
    )
    parser.add_argument(
        "--git_user",
        default=subprocess.getoutput(["git config user.name"]),
        help="Username to set as the git user name in container.",
    )
    parser.add_argument(
        "--git_email",
        default=subprocess.getoutput(["git config user.email"]),
        help="Email to set as git email in container.",
    )

    parser.add_argument(
        "--image_prefix", default="roller", help="Prefix for image names."
    )
    parser.add_argument(
        "--container_prefix", default="dev", help="Prefix for container name."
    )

    parser.add_argument(
        "--user_dockerfile",
        default="docker/user-image/Dockerfile",
        help="User docker file.",
    )

    parser.add_argument(
        "--privileged",
        action="store_true",
        help="Whether to run the container in privileged mode.",
    )

    parser.add_argument(
        "--versions",
        default=dockerfiles.keys(),
        choices=dockerfiles.keys(),
        nargs="+",
        help="Which container versions to build.",
    )

    parser.add_argument(
        "--context",
        default="./",
        help="What context to use when building the containers.",
    )
    parser.add_argument(
        "--volumes",
        default=[gitTop().as_posix() + ":/data"],
        action="append",
        help="What volumes to mount to the container. e.g. /source:/dest",
    )

    parser.add_argument(
        "--delete_existing_container",
        action="store_true",
        help="Whether to remove an already running container before launching a new one.",
    )

    parser.add_argument(
        "--interactive_args",
        default="-dt",
        dest="interactive_arg",
        help="Interactive args to pass to docker. E.g. -ti or -dt",
    )
    parser.add_argument(
        "--docker_gpu_args",
        default=None,
        action="append",
        help="Used to override default docker GPU args.",
    )
    parser.add_argument(
        "--base_build_extra_args",
        default=[],
        action="append",
        help="Used to pass extra args to the docker build of the base image.",
    )
    parser.add_argument(
        "--user_build_extra_args",
        default=[],
        action="append",
        help="Used to pass extra args to the docker build of the user image.",
    )
    parser.add_argument(
        "--run_extra_args",
        default=[],
        action="append",
        help="Used to pass extra args to the docker run of the user container.",
    )

    parser.add_argument(
        "--skip_base_build",
        action="store_true",
        help="Whether to skip the base image build.",
    )
    parser.add_argument(
        "--skip_user_build",
        action="store_true",
        help="Whether to skip the user image build.",
    )
    parser.add_argument(
        "--skip_container_run",
        action="store_true",
        help="Whether to skip launching the container.",
    )

    parser.add_argument("--amdgpu-url", help="amdgpu-install DEB URL")

    parser.add_argument("--amdgpu-build-number", help="amdgpu-install build number")

    parser.add_argument("--amdgpu-build-uri", help="amdgpu-install build uri")

    parser.add_argument(
        "--pull",
        dest="base_build_extra_args",
        action="append_const",
        const="--pull",
        help="Whether to try to pull newer base images",
    )

    parser.add_argument("command", nargs=argparse.REMAINDER)

    return parser.parse_args()


def buildx_available():
    try:
        subprocess.run(
            ["docker", "buildx", "version"],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
            check=True,
        )
        return True
    except subprocess.CalledProcessError:
        return False


def buildImage(dockerfile, name, context="./", *args):
    build_command = ["docker", "build"]
    if buildx_available():
        build_command = ["docker", "buildx", "build"]

    printAndRun(build_command + ["-t", name, "-f", dockerfile] + list(args) + [context])


def get_render_group():
    try:
        return grp.getgrnam("render").gr_gid
    except KeyError:
        return None


def get_render_group_args():
    render_group = get_render_group()
    if render_group is None:
        return []

    return ["-e", "_RENDER_GID=" + str(render_group)]


def get_run_args(args, image_name, container_name):
    volume_args = sum([["-v", volume] for volume in args.volumes], [])

    env_args = [
        "-e",
        "_USER=" + args.user,
        "-e",
        "_UID=" + args.uid,
        "-e",
        "_GIT_USER=" + args.git_user,
        "-e",
        "_GIT_EMAIL=" + args.git_email,
    ]

    env_args += get_render_group_args()
    if args.privileged:
        env_args += ["--privileged"]

    run_args = ["docker", "run"]
    run_args += env_args
    run_args += volume_args
    run_args += [args.interactive_arg]
    run_args += args.docker_gpu_args
    run_args += ["--name", container_name, image_name + ":latest"]
    run_args += args.run_extra_args
    run_args += args.command

    return run_args


def get_amdgpu_args(args):
    build_args = []
    for var in [
        "ROCROLLER_AMDGPU_URL",
        "ROCROLLER_AMDGPU_BUILD_NUMBER",
        "ROCROLLER_AMDGPU_BUILD_URI",
    ]:
        prefix = "ROCROLLER_"
        arg = var[len(prefix) :].lower()
        if var in os.environ:
            build_args.append(var + "=" + os.environ[var])
        if getattr(args, arg, None) is not None:
            if var in os.environ:
                raise RuntimeError(
                    f"{var} should not be specified both via environment variable and command-line argument."
                )
            build_args.append(var + "=" + getattr(args, arg))

    return build_args


def make_build_args(build):
    rv = []
    for arg in build:
        rv.extend(["--build-arg", arg])
    return rv


def main():
    final_output = ""
    args = parseArgs()

    if args.docker_gpu_args is None:
        args.docker_gpu_args = [
            "--device=/dev/kfd",
            "--device=/dev/dri",
            "--security-opt=seccomp=unconfined",
            "--net=host",
        ]

    if not args.git_user or not args.git_email:
        final_output += (
            "WARNING: Git configuration was not detectable, use 'git config' to set "
            "user.name and user.email to have them automatically passed in.\n"
        )

    if len(args.command) > 0:
        args.command = ["bash", "--login", "-c", " ".join(args.command)]

    build_args = get_amdgpu_args(args)
    print(build_args)

    for version in args.versions:
        base_image_name = "-".join([args.image_prefix, version])
        user_image_name = "-".join([base_image_name, "user"])
        container_name = "_".join([args.user, args.container_prefix, version])

        if args.delete_existing_container:
            printAndRun(["docker", "rm", "-f", container_name])

        if not args.skip_base_build:
            buildImage(
                dockerfiles[version],
                base_image_name,
                args.context,
                *make_build_args(build_args),
                *[
                    arg
                    for extra_arg in args.base_build_extra_args
                    for arg in extra_arg.split()
                ],
            )

        if not args.skip_user_build:
            buildImage(
                args.user_dockerfile,
                user_image_name,
                args.context,
                *make_build_args(
                    [
                        "base_image=" + base_image_name,
                        "_UID=" + args.uid,
                        "_GID=" + str(get_render_group()),
                    ]
                ),
                *[
                    arg
                    for extra_arg in args.user_build_extra_args
                    for arg in extra_arg.split()
                ],
            )

        if not args.skip_container_run:
            run_args = get_run_args(args, user_image_name, container_name)
            printAndRun(run_args)
            final_output += (
                " ".join(["docker exec -ti -u", args.user, container_name, "bash"])
                + "\n"
            )

    if final_output:
        print("Enter these containers with the following commands:")
        print(final_output)


def printAndRun(run_args):
    print(run_args)
    print(" ".join(run_args))
    subprocess.check_call(run_args)


if __name__ == "__main__":
    main()