#!/usr/bin/env python3

# Copyright 2023 The IREE Authors
#
# Licensed under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""Configure script to setup VSCode DevContainer Environment."""

import json
import os
import pathlib
import shlex
import subprocess
import sys

from importlib.machinery import SourceFileLoader
from shutil import which

CURRENT_DIR = pathlib.Path(__file__).resolve().parent

DOCKER_IMAGE_SHORTNAME_DICT = {
    "base": "base-bleeding-edge",
    "nvidia": "nvidia-bleeding-edge",
}


def run_shell(cmd):
    cmd = shlex.split(cmd)
    return subprocess.run(cmd, check=True, text=True, stdout=subprocess.PIPE).stdout


def is_nvidia_gpu_available():
    """This function verifies NVIDIA Docker runtime is installed and
    available. It also checks NVIDIA GPU are available and the driver
    is installed."""

    command_str = f"'{which('docker')}' info --format '{{{{json .}}}}'"
    data = json.loads(run_shell(command_str))

    if "nvidia" not in data["Runtimes"].keys():
        return (
            False,
            "NVIDIA Docker Runtime is not available. Please see: "
            "https://developer.nvidia.com/nvidia-container-runtime for additional "
            "information.",
        )

    nvidia_smi_executable = which("nvidia-smi")
    if nvidia_smi_executable is None:
        return False, "NVIDIA Driver is not installed or not in the user path."

    command_str = f"{nvidia_smi_executable} --query-gpu=gpu_name --format=csv"

    if run_shell(command_str).splitlines()[1:]:  # Skip header
        return True, None
    else:
        return False, "No NVIDIA GPU was found on this system."


def _get_input(question, default_answer):
    try:
        answer = input(f"\n- {question} ")
    except EOFError:
        answer = default_answer

    return (answer or default_answer).strip().lower()


def get_input(question, default_answer="", accepted_answers=None):
    if accepted_answers is None:
        raise RuntimeError("Argument `accepted_answers` is None.")

    accepted_answers = [x.strip().lower() for x in accepted_answers]

    while True:
        answer = _get_input(question, default_answer)
        if answer not in accepted_answers:
            print(
                f"\tERROR: Unsupported answer received: {answer}."
                f"Expected: {accepted_answers}"
            )
            continue
        break

    return answer


def get_input_path(question):
    while True:
        answer = _get_input(question, default_answer="")
        if answer != "" and not os.path.isdir(answer):
            print(f"\tERROR: Received path does not exist: `{answer}`")
            continue
        break

    return answer


if __name__ == "__main__":
    # ------------------------------------------------------------------------- #
    #                         Environment Verifications                         #
    # ------------------------------------------------------------------------- #

    docker_executable = which("docker")
    if docker_executable is None:
        raise RuntimeError(
            "Docker is not installed or in the user path. Please refer to: "
            "https://docs.docker.com/desktop/"
        )

    try:
        run_shell(f"'{docker_executable}' compose version")
    except subprocess.CalledProcessError as e:
        raise RuntimeError(
            "Docker Compose must be installed in order to use IREE VS Code "
            "Development Container. Please refer to: "
            "https://docs.docker.com/compose/"
        ) from e

    docker_image_key = "base"

    # ------------------------------------------------------------------------- #
    #                               Mandatory Steps                             #
    # ------------------------------------------------------------------------- #

    # STEP 1: Verify the user doesn't have a pre-existing `docker-compose.yml`.
    #         If yes, confirm they want to overwrite it.
    if os.path.isfile(CURRENT_DIR / "docker-compose.yml"):
        if (
            get_input(
                "A `docker-compose.yml` already exists. Are you certain you want to overwrite it [y/N]?",
                default_answer="n",
                accepted_answers=["y", "n"],
            )
            == "n"
        ):
            sys.exit(1)

    # STEP 2: Read the template configuration file
    with open(CURRENT_DIR / "docker-compose.base.yml") as f:
        docker_config = json.load(f)

    # STEP 3: Prebuilt vs Locally Built Containers
    use_official_img = (
        get_input(
            "Do you wish to use the official prebuild development containers [Y/n]?",
            default_answer="y",
            accepted_answers=["y", "n"],
        )
        == "y"
    )

    # ------------------------------------------------------------------------- #
    #                 Optional Mounting Points - Build & Cache                  #
    # ------------------------------------------------------------------------- #

    # STEP 4 [OPTIONAL]: Does the user want to mount a directory for CCACHE ?
    ccache_mount_cache = get_input_path(
        "[OPTIONAL] Enter the the directory path to mount as CCache Cache.\n"
        "  It will be mounted in the container under: `/home/iree/.cache/ccache` [Default: None]:"
    )
    if ccache_mount_cache:
        docker_config["services"]["iree-dev"]["volumes"].append(
            f"{ccache_mount_cache}:/home/iree/.cache/ccache:cached"
        )

    # STEP 5 [OPTIONAL]: Does the user want to mount a directory for BUILD ?
    build_mount_cache = get_input_path(
        "[OPTIONAL] Enter the the directory path to mount as Build Directory.\n"
        "  It will mounted in the container under: `/home/iree/build` [Default: None]:"
    )
    if build_mount_cache:
        docker_config["services"]["iree-dev"]["volumes"].append(
            f"{build_mount_cache}:/home/iree/build:cached"
        )

    # ------------------------------------------------------------------------- #
    #                Optional Deep Learning Accelerator Support                 #
    # ------------------------------------------------------------------------- #

    # STEP 6 [OPTIONAL]: Does the user want to use NVIDIA GPUs ?
    nvidia_gpu_available, err_msg = is_nvidia_gpu_available()

    if nvidia_gpu_available:
        if (
            get_input(
                "[OPTIONAL] Do you wish to use NVIDIA GPU [y/N]?",
                default_answer="n",
                accepted_answers=["y", "n"],
            )
            == "y"
        ):
            docker_image_key = "nvidia"

    else:
        print(f"[INFO] NVIDIA GPUs are not available for use: {err_msg}")

    if docker_image_key != "nvidia":
        del docker_config["services"]["iree-dev"]["deploy"]

    # ------------------------------------------------------------------------- #
    #            Setting the right docker image / container to build            #
    # ------------------------------------------------------------------------- #

    docker_img_shortname = DOCKER_IMAGE_SHORTNAME_DICT[docker_image_key]

    if use_official_img:
        del docker_config["services"]["iree-dev"]["build"]

        docker_iree_registry = SourceFileLoader(
            "docker_iree_registry",
            str(CURRENT_DIR / "../build_tools/docker/get_image_name.py"),
        ).load_module()

        image_name = docker_iree_registry.find_image_by_name(docker_img_shortname)
        docker_config["services"]["iree-dev"]["image"] = image_name

    else:
        del docker_config["services"]["iree-dev"]["image"]

        dockerfile = f"build_tools/docker/dockerfiles/{docker_img_shortname}.Dockerfile"
        docker_config["services"]["iree-dev"]["build"]["dockerfile"] = dockerfile

        if not os.path.isfile(CURRENT_DIR / ".." / dockerfile):
            raise FileNotFoundError(f"The file `{dockerfile}` does not exist.")

    docker_compose_filepath = os.path.join(CURRENT_DIR, "docker-compose.yml")
    with open(docker_compose_filepath, "w") as f:
        json.dump(docker_config, f, indent=2)

    # ------------------------------------------------------------------------- #
    #                                  SUCCESS                                  #
    # ------------------------------------------------------------------------- #

    print("\n" + "=" * 80)
    print(f"\nSuccess! Wrote Docker Compose file to `{docker_compose_filepath}`.")