# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.

import argparse
import importlib.util as import_util
import os
import pathlib
import subprocess
import sys
from typing import List, Optional, Sequence, Union

VENV_NAME = ".venv"
CWD = pathlib.PurePath(os.getcwd())


class VenvError(Exception):
    pass


def parse_args(argv: Sequence[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--requirements",
        action="append",
        default=[],
        help="Install additional dependencies into the virtual environment.",
    )

    parser.add_argument(
        "--toml",
        action="store",
        default=None,
        help="Install additional dependencies from sources like `pyproject.toml` into the virtual environment.",
    )
    parser.add_argument(
        "--extras",
        action="append",
        default=[],
        help="Install specific package groups from `pyproject.toml` into the virtual environment.",
    )

    parser.add_argument(
        "--git-ignore",
        action="store_true",
        default=False,
        help="Add .gitignore to the newly created virtual environment.",
    )
    parser.add_argument(
        "--name",
        default=VENV_NAME,
        type=str,
        help="Name of the virtual environment.",
        metavar="NAME",
        action="store",
    )
    return parser.parse_args(argv)


def is_installed(module: str) -> bool:
    return import_util.find_spec(module) is not None


def file_exists(path: Union[str, pathlib.PurePath]) -> bool:
    return os.path.exists(path)


def venv_exists(name: str) -> bool:
    return os.path.exists(CWD / name) and file_exists(get_venv_path(name))


def run_process(args: Sequence[str], error_message: str) -> None:
    try:
        print("Running: " + " ".join(args))
        subprocess.run(args, cwd=os.getcwd(), check=True)
    except subprocess.CalledProcessError:
        raise VenvError(error_message)


def get_venv_path(name: str) -> str:
    # See `venv` doc here for more details on binary location:
    # https://docs.python.org/3/library/venv.html#creating-virtual-environments
    if sys.platform == "win32":
        return os.fspath(CWD / name / "Scripts" / "python.exe")
    else:
        return os.fspath(CWD / name / "bin" / "python")


def install_requirements(venv_path: str, requirements: List[str]) -> None:
    if not requirements:
        return

    print(f"VENV_INSTALLING_REQUIREMENTS: {requirements}")
    args = []
    for requirement in requirements:
        args += ["-r", requirement]
    run_process(
        [venv_path, "-m", "pip", "install"] + args,
        "CREATE_VENV.PIP_FAILED_INSTALL_REQUIREMENTS",
    )
    print("CREATE_VENV.PIP_INSTALLED_REQUIREMENTS")


def install_toml(venv_path: str, extras: List[str]) -> None:
    args = "." if len(extras) == 0 else f".[{','.join(extras)}]"
    run_process(
        [venv_path, "-m", "pip", "install", "-e", args],
        "CREATE_VENV.PIP_FAILED_INSTALL_PYPROJECT",
    )
    print("CREATE_VENV.PIP_INSTALLED_PYPROJECT")


def upgrade_pip(venv_path: str) -> None:
    run_process(
        [venv_path, "-m", "pip", "install", "--upgrade", "pip"],
        "CREATE_VENV.PIP_UPGRADE_FAILED",
    )


def add_gitignore(name: str) -> None:
    git_ignore = CWD / name / ".gitignore"
    if not file_exists(git_ignore):
        print("Creating: " + os.fspath(git_ignore))
        with open(git_ignore, "w") as f:
            f.write("*")


def main(argv: Optional[Sequence[str]] = None) -> None:
    if argv is None:
        argv = []
    args = parse_args(argv)

    if not is_installed("venv"):
        raise VenvError("CREATE_VENV.VENV_NOT_FOUND")

    pip_installed = is_installed("pip")
    deps_needed = args.requirements or args.extras or args.toml
    if deps_needed and not pip_installed:
        raise VenvError("CREATE_VENV.PIP_NOT_FOUND")

    if venv_exists(args.name):
        venv_path = get_venv_path(args.name)
        print(f"EXISTING_VENV:{venv_path}")
    else:
        run_process(
            [sys.executable, "-m", "venv", args.name],
            "CREATE_VENV.VENV_FAILED_CREATION",
        )
        venv_path = get_venv_path(args.name)
        print(f"CREATED_VENV:{venv_path}")
        if args.git_ignore:
            add_gitignore(args.name)

    if pip_installed:
        upgrade_pip(venv_path)

    if args.toml:
        print(f"VENV_INSTALLING_PYPROJECT: {args.toml}")
        install_toml(venv_path, args.extras)

    if args.requirements:
        print(f"VENV_INSTALLING_REQUIREMENTS: {args.requirements}")
        install_requirements(venv_path, args.requirements)


if __name__ == "__main__":
    main(sys.argv[1:])