"""Create Package variants for PyPI distribution."""

import argparse
import os

from test_utils import PY_PACKAGE, ROOT

IN_PATH = os.path.join(PY_PACKAGE, "pyproject.toml.in")
OUT_PATH = os.path.join(PY_PACKAGE, "pyproject.toml")

WHL_CPU = """
[tool.hatch.build.targets.wheel]
packages = ["xgboost/"]
"""

CHOICES = ["default", "cpu", "manylinux2014"]


def copyfile(src: str, dst: str) -> None:
    with open(src, "rb") as fd:
        content = fd.read()
    with open(dst, "wb") as fd:
        fd.write(content)


def make_pyproject(variant: str) -> None:
    assert variant in CHOICES

    with open(IN_PATH) as fd:
        pyproject = fd.read()

    readme_dft = os.path.join(PY_PACKAGE, "README.dft.rst")
    readme_cpu = os.path.join(PY_PACKAGE, "README.cpu.rst")
    readme = os.path.join(PY_PACKAGE, "README.rst")
    if variant == "cpu":
        pyproject = (
            pyproject.replace("{{ name }}", "xgboost-cpu")
            .replace("{{ wheel }}", WHL_CPU)
            .replace("{{ nccl }}", "")
        )
        copyfile(readme_cpu, readme)
    elif variant == "manylinux2014":
        pyproject = (
            pyproject.replace("{{ name }}", "xgboost")
            .replace("{{ wheel }}", "")
            .replace("{{ nccl }}", "")
        )
        copyfile(readme_dft, readme)
    else:
        pyproject = (
            pyproject.replace("{{ name }}", "xgboost")
            .replace("{{ wheel }}", "")
            .replace(
                "{{ nccl }}",
                """    \"nvidia-nccl-cu12 ; platform_system == 'Linux' and platform_machine != 'aarch64'\",""",
            )
        )
        copyfile(readme_dft, readme)
    pyproject = (
        f"# Generated by `{os.path.basename(__file__)}`, don't edit.\n" + pyproject
    )

    with open(OUT_PATH, "w") as fd:
        fd.write(pyproject)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--variant",
        type=str,
        choices=CHOICES,
        default="default",
    )
    args = parser.parse_args()
    make_pyproject(args.variant)
