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

import argparse
import os
import pathlib
import subprocess
import sys
from typing import Optional, Sequence, Union

CONDA_ENV_NAME = ".conda"
CWD = pathlib.Path.cwd()


class VenvError(Exception):
    pass


def parse_args(argv: Sequence[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--python",
        action="store",
        help="Python version to install in the virtual environment.",
        default=f"{sys.version_info.major}.{sys.version_info.minor}",
    )
    parser.add_argument(
        "--install",
        action="store_true",
        default=False,
        help="Install packages 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=CONDA_ENV_NAME,
        type=str,
        help="Name of the virtual environment.",
        metavar="NAME",
        action="store",
    )
    return parser.parse_args(argv)


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


def conda_env_exists(name: Union[str, pathlib.PurePath]) -> bool:
    return os.path.exists(CWD / 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_conda_env_path(name: str) -> str:
    return os.fspath(CWD / name)


def install_packages(env_path: str) -> None:
    yml = os.fspath(CWD / "environment.yml")
    if file_exists(yml):
        print(f"CONDA_INSTALLING_YML: {yml}")
        run_process(
            [
                sys.executable,
                "-m",
                "conda",
                "env",
                "update",
                "--prefix",
                env_path,
                "--file",
                yml,
            ],
            "CREATE_CONDA.FAILED_INSTALL_YML",
        )
        print("CREATE_CONDA.INSTALLED_YML")


def add_gitignore(name: str) -> None:
    git_ignore = os.fspath(CWD / name / ".gitignore")
    if not file_exists(git_ignore):
        print(f"Creating: {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 conda_env_exists(args.name):
        env_path = get_conda_env_path(args.name)
        print(f"EXISTING_CONDA_ENV:{env_path}")
    else:
        run_process(
            [
                sys.executable,
                "-m",
                "conda",
                "create",
                "--yes",
                "--prefix",
                args.name,
                f"python={args.python}",
            ],
            "CREATE_CONDA.ENV_FAILED_CREATION",
        )
        env_path = get_conda_env_path(args.name)
        print(f"CREATED_CONDA_ENV:{env_path}")
        if args.git_ignore:
            add_gitignore(args.name)

    if args.install:
        install_packages(env_path)


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