#!/usr/bin/env python
"""Generate requirements/*.txt files from pyproject.toml.

Also builda a conda environment.yml
"""

import sys
import re
from pathlib import Path

try:  # standard module since Python 3.11
    import tomllib as toml
except ImportError:
    try:  # available for older Python via pip
        import tomli as toml
    except ImportError:
        sys.exit("Please install `tomli` first: `pip install tomli`")

script_pth = Path(__file__)
repo_dir = script_pth.parent.parent
script_relpth = script_pth.relative_to(repo_dir)
header = [
    f"# Generated via {script_relpth.as_posix()} and pre-commit hook.",
    "# Do not edit this file; modify pyproject.toml instead.",
]


def generate_requirement_file(name: str, req_list: list[str]) -> None:
    req_fname = repo_dir / "requirements" / f"{name}.txt"
    req_fname.write_text("\n".join(header + req_list) + "\n")


def generate_environment_yml(req_sections: dict[str, list[str]]) -> None:
    # Some packages in conda-forge have different names than they do on pypi
    # Also, some packages have a 'base' flavoured conda package which
    # doesn't install all their optional dependencies.
    rename_idx = {
        'build': 'python-build',
        'kaleido': 'python-kaleido',
        'sphinx_design': 'sphinx-design',
        'astropy': 'astropy-base',
        'matplotlib': 'matplotlib-base',
    }
    lines = ["name: skimage-dev", "channels:", "  - conda-forge", "dependencies:"]
    for section in req_sections:
        lines.append(f"  # {section}")
        for dep in req_sections[section]:
            # Remove optional specifiers such as `[parallel]`
            dep = re.sub('\\[.*?\\]', '', dep)

            # Remove platform specifiers such as `; sys_platform != "emscripten"`
            dep = re.sub('; .*', '', dep)

            pkgname = re.split('[>=]', dep)[0]
            dep = dep.replace(pkgname, rename_idx.get(pkgname, pkgname))
            if dep == "scikit-image":
                continue

            lines.append(f"  - {dep}")

    with open("environment.yml", "w") as f:
        f.writelines(f"{line}\n" for line in lines)


def expand_dependencies(
    dep_list: list[str],
    optional_dict: dict[str, list[str]],
    package_name: str = 'scikit-image',
) -> list[str]:
    """Explode dependencies with optional extras into a flat list.

    If `scikit-image[optional_group]` is used as a dependency itself, replace
    with the actual dependencies of `optional_group`.
    """
    exploded = []
    for dep in dep_list:
        if dep.startswith(package_name):
            extras = dep.split('[')[1].rstrip(']').split(',')
            exploded.extend(
                expand_dependencies(
                    optional_dict[extras[0]], optional_dict, package_name
                )
            )
        else:
            exploded.append(dep)
    return exploded


def main() -> None:
    pyproject = toml.loads((repo_dir / "pyproject.toml").read_text())

    generate_requirement_file("default", pyproject["project"]["dependencies"])

    for key, opt_list in pyproject["project"]["optional-dependencies"].items():
        generate_requirement_file(
            key,
            expand_dependencies(
                opt_list,
                pyproject["project"]["optional-dependencies"],
                package_name='scikit-image',
            ),
        )

    generate_environment_yml(
        {
            **{'base': pyproject["project"]["dependencies"]},
            **pyproject["project"]["optional-dependencies"],
        }
    )


if __name__ == "__main__":
    main()
