"""Select the image builds and GPU verifications for a scoped release tag."""

from __future__ import annotations

import argparse
import json
import re
from dataclasses import dataclass
from pathlib import Path


GPU_GROUPS = {
    "nvidia-a100": ("Nvidia-A100",),
    "cambricon-mlu370": ("Cambricon_mlu-370-x4", "Cambricon_mlu-370-x8"),
    "metax-c500": ("MetaX_c-500",),
    "vastai-va16": ("Vastai_va16",),
}
RELEASE_TAG = re.compile(
    r"v[0-9]+\.[0-9]+\.[0-9]+-(nvidia-a100|cambricon-mlu370|metax-c500|vastai-va16|all)\Z"
)


@dataclass(frozen=True)
class ReleasePlan:
    target: str
    image_targets: tuple[str, ...]
    gpu_types: tuple[str, ...]
    verification_pairs: tuple[tuple[str, str], ...]


def plan_for_tag(tag: str) -> ReleasePlan:
    match = RELEASE_TAG.fullmatch(tag)
    if match is None:
        raise ValueError(
            "release tag must be vX.Y.Z-<target>, where target is "
            "nvidia-a100, cambricon-mlu370, metax-c500, vastai-va16, or all"
        )
    target = match.group(1)
    image_targets = tuple(GPU_GROUPS) if target == "all" else (target,)
    verification_pairs = tuple(
        (gpu, group) for group in image_targets for gpu in GPU_GROUPS[group]
    )
    gpu_types = tuple(gpu for gpu, _ in verification_pairs)
    return ReleasePlan(target, image_targets, gpu_types, verification_pairs)


def write_gitea_env(plan: ReleasePlan, path: Path) -> None:
    with path.open("a", encoding="utf-8") as env_file:
        env_file.write(f"RELEASE_TARGET={plan.target}\n")
        env_file.write(f"IMAGE_TARGETS={' '.join(plan.image_targets)}\n")
        env_file.write(f"GPU_TYPES={' '.join(plan.gpu_types)}\n")
        env_file.write(
            "VERIFY_PAIRS="
            + " ".join(f"{gpu}:{group}" for gpu, group in plan.verification_pairs)
            + "\n"
        )


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("tag", help="scoped release tag, such as v0.3.0-metax-c500")
    parser.add_argument("--gitea-env", type=Path, help="append the plan to GITEA_ENV")
    args = parser.parse_args()
    try:
        plan = plan_for_tag(args.tag)
    except ValueError as exc:
        parser.error(str(exc))
    if args.gitea_env:
        write_gitea_env(plan, args.gitea_env)
    print(json.dumps(plan.__dict__, ensure_ascii=False))


if __name__ == "__main__":
    main()