#!/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
from argh import arg  # type: ignore
from impl.common import CROSVM_ROOT, run_main, cmd, chdir, quoted
import sys
import os

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

try:
    docker = cmd(os.environ.get("DOCKER", "docker"))
except ValueError:
    docker = cmd("podman")

is_podman = docker.executable.name == "podman"

# Enable interactive mode when running in an interactive terminal.
TTY_ARGS = "--interactive --tty" if sys.stdin.isatty() else None


DOCKER_ARGS = [
    TTY_ARGS,
    # Podman will not share devices when `--privileged` is specified
    "--privileged" if not is_podman else None,
    # 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}",
]


def container_revision(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]


@arg("command", nargs=argparse.REMAINDER)
def main(command: tuple[str, ...], stop: bool = False, hermetic: bool = False):
    chdir(CROSVM_ROOT)
    container_id = docker(f"ps -q -f name={CONTAINER_NAME}").stdout()

    # 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", *DOCKER_ARGS, *quoted_cmd).fg()
    else:
        if container_id and container_revision(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}", *DOCKER_ARGS).stdout()
            print(f"Started dev-container ({container_id}).")
        else:
            print(f"Using existing dev-container instance ({container_id}).")

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


if __name__ == "__main__":
    run_main(main)