#!/usr/bin/env python3
# Copyright 2021 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
#
# Usage:
#
# To get an interactive shell for development:
#   ./tools/dev_container
#
# To run a command in the container, e.g. to run presubmits:
#   ./tools/dev_container ./tools/presubmit
#
# The state of the container (including build artifacts) are preserved between
# calls. To stop the container call:
#   ./tools/dev_container --stop
#
# The dev container can also be called with a fresh container for each call that
# is cleaned up afterwards (e.g. when run by Kokoro):
#
#   ./tools/dev_container --hermetic CMD

import argparse
import getpass
import shutil
from typing import Tuple
from impl.common import CROSVM_ROOT, cmd, chdir, quoted
import sys

CONTAINER_NAME = f"crosvm_dev_{getpass.getuser()}"
IMAGE_VERSION = (CROSVM_ROOT / "tools/impl/dev_container/version").read_text().strip()

DOCKER_ARGS = [
    # Share crosvm source
    f"--volume {quoted(CROSVM_ROOT)}:/workspace:rw",
    # Share devices and syslog
    "--device /dev/kvm",
    "--volume /dev/log:/dev/log",
    "--device /dev/net/tun",
    "--device /dev/vhost-net",
    "--device /dev/vhost-vsock",
    # Use tmpfs in the container for faster performance.
    "--mount type=tmpfs,destination=/tmp",
    # For plugin process jail
    "--mount type=tmpfs,destination=/var/empty",
    f"gcr.io/crosvm-packages/crosvm_dev:{IMAGE_VERSION}",
]

PODMAN_IS_DEFAULT = shutil.which("docker") == None


def container_revision(docker: cmd, container_id: str):
    image = docker("container inspect -f {{.Config.Image}}", container_id).stdout()
    parts = image.split(":")
    assert len(parts) == 2, f"Invalid image name {image}"
    return parts[1]


def main(
    command: Tuple[str, ...],
    stop: bool = False,
    hermetic: bool = False,
    interactive: bool = False,
    podman: bool = PODMAN_IS_DEFAULT,
):
    chdir(CROSVM_ROOT)

    docker = cmd("podman" if podman else "docker")
    docker_args = [
        # Podman will not share devices when `--privileged` is specified
        "--privileged" if not podman else None,
        *DOCKER_ARGS,
    ]

    if podman:
        print("WARNING: Running dev_container with podman is not fully supported.")
        print("Some crosvm tests require privileges podman cannot provide and may fail.")
        print()

    container_id = docker(f"ps -a -q -f name={CONTAINER_NAME}").stdout()

    # If a command is provided run non-interactive unless explicitly asked for.
    tty_args = []
    if not command or interactive:
        if not sys.stdin.isatty():
            raise Exception("Trying to run an interactive session in a non-interactive terminal.")
        tty_args = ["--interactive", "--tty"]

    # Start an interactive shell by default
    if not command:
        command = ("/bin/bash",)

    quoted_cmd = list(map(quoted, command))

    if stop:
        if container_id:
            print(f"Stopping dev-container {container_id}.")
            docker("rm -f", container_id).fg(quiet=True)
        else:
            print(f"Dev-container is not running.")
        return

    if hermetic:
        docker(f"run --rm", *tty_args, *docker_args, *quoted_cmd).fg()
    else:
        if container_id and container_revision(docker, container_id) != IMAGE_VERSION:
            print(f"New image is available. Stopping old container ({container_id}).")
            docker("rm -f", container_id).fg(quiet=True)
            container_id = None

        if not container_id:
            container_id = docker(
                f"run --detach --name {CONTAINER_NAME}", *tty_args, *docker_args
            ).stdout()
            print(f"Started dev-container ({container_id}).")
        else:
            is_running = docker(f"ps -q -f name={CONTAINER_NAME}").stdout()
            if not is_running:
                # The container may have been stopped (e.g. by a reboot)
                print(f"Starting existing dev-container ({container_id}).")
                docker("start", container_id).fg()
            print(f"Using existing dev-container instance ({container_id}).")

        docker("exec", *tty_args, container_id, *quoted_cmd).fg()


if __name__ == "__main__":
    # Note: Since the dev_container script is run outside of the container, we do not want to
    # use common.run_main/argh here, so we can use this script without installing any
    # dependencies. All that's needed is python >= 3.8.
    parser = argparse.ArgumentParser(description="Runs the crosvm dev container.")
    parser.add_argument("--stop", action="store_true", help="Stop the container if it's running.")
    parser.add_argument(
        "--hermetic", action="store_true", help="Do not re-use an existing container."
    )
    parser.add_argument(
        "--interactive", "-i", action="store_true", help="Run the command interactively."
    )
    parser.add_argument(
        "--podman",
        action="store_true",
        default=PODMAN_IS_DEFAULT,
        help="Use podman instead of docker.",
    )
    parser.add_argument(
        "command", nargs=argparse.REMAINDER, help="The command to execute in the container."
    )
    args = parser.parse_args()
    main(
        args.command,
        stop=args.stop,
        hermetic=args.hermetic,
        interactive=args.interactive,
        podman=args.podman,
    )