#!/usr/bin/python3

from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser, FileType
from os import devnull
from pathlib import Path
from sys import exit, stderr

from tomli import load


def parse_args():
    p = ArgumentParser(
        description=("Write hatchling environments to requirements files"),
        formatter_class=ArgumentDefaultsHelpFormatter,
        epilog="""The output directory will files of the form env.{foo,bar}.txt
        for any environments (foo, bar) defined in envs in hatch.toml, if
        present, or in tool.hatch.envs in pyproject.toml.""",
    )
    p.add_argument(
        "--pyproject",
        metavar="FILE",
        type=FileType("rb"),
        default=Path("pyproject.toml"),
        help="Path to pyproject.toml file that uses hatchling",
    )
    p.add_argument(
        "--hatch",
        metavar="FILE",
        type=FileType("rb"),
        default=Path("hatch.toml"),
        help="Path to optional hatch.toml file",
    )
    p.add_argument(
        "--dir",
        metavar="DIR",
        type=Path,
        default="_req",
        help="Directory to contain requirements files",
    )
    p.add_argument(
        "--verbose",
        "-v",
        action="store_const",
        const=stderr,
        default=Path(devnull),
        help="Print details about what is happening",
    )
    args = p.parse_args()
    # We don’t use a with statement to close this file, because if it is
    # stderr, we don’t want to close it before reporting a traceback.
    args.verbose = ensure_open(args.verbose, "w")
    return args


def ensure_open(path_or_file, mode):
    if isinstance(path_or_file, Path):
        path_or_file = path_or_file.open(mode)
    return path_or_file


def dump(output, dependencies, v):
    with output.open("w") as requirements:
        requirements.writelines(
            f"{dependency}\n" for dependency in dependencies
        )
    n = len(dependencies)
    print(
        f"Wrote {output.name} with {n} dependenc{'y' if n == 1 else 'ies'}",
        file=v,
    )


def resolve_environment(env_name, all_envs, project):
    # https://hatch.pypa.io/latest/config/environment/
    env = all_envs[env_name]
    deps = []
    try:
        # If "dependencies" is present, it overrides templates anyway.
        deps = env["dependencies"]
    except KeyError:
        if not env.get("detached", False):
            # Inherit from templates
            template = env.get("template", "default")
            if template != env_name:
                deps = resolve_environment(template, all_envs, project)
    for extra in env.get("features", []):
        deps.extend(project["optional-dependencies"][extra])
    deps.extend(env.get("extra-dependencies", []))
    return deps


def main():
    args = parse_args()
    v = args.verbose

    with ensure_open(args.pyproject, "rb") as pyproject_file:
        pyproject = load(pyproject_file)
    try:
        with ensure_open(args.hatch, "rb") as hatch_file:
            hatch = load(hatch_file)
    except FileNotFoundError:
        hatch = {}

    backend = pyproject.get("build-system", {}).get("build-backend")
    if backend != "hatchling.build":
        return "build-system.build-backend is not hatchling.build"
    try:
        project = pyproject["project"]
    except KeyError:
        return "no “project” section is present"

    args.dir.mkdir(parents=True, exist_ok=True)
    print(f"Created {args.dir.resolve()}", file=v)

    try:
        envs = hatch["envs"]
    except KeyError:
        envs = pyproject.get("tool", {}).get("hatch", {}).get("envs", {})
    for name in envs:
        requirements = args.dir / f"env.{name}.txt"
        dependencies = resolve_environment(name, envs, project)
        dump(requirements, dependencies, v)


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