computer-playbook/cli/deploy/development/init.py

# cli/deploy/development/init.py
from __future__ import annotations

import argparse
import json
import os
from typing import Any, Dict

from .common import make_compose, resolve_deploy_ids_for_apps
from .inventory import (
    DevInventorySpec,
    build_dev_inventory,
    filter_plan_to_variant,
    plan_dev_inventory_matrix,
)
from .storage import detect_storage_constrained
from ...meta.runtime import detect_runtime


def add_parser(sub: argparse._SubParsersAction) -> None:
    p = sub.add_parser(
        "init",
        help="Create development inventory inside the infinito container.",
    )
    p.add_argument(
        "--distro",
        default=os.environ.get("INFINITO_DISTRO", "arch"),
        choices=["arch", "debian", "ubuntu", "fedora", "centos"],
        help="Target distro (compose env INFINITO_DISTRO).",
    )

    p.add_argument(
        "--inventory-dir",
        default=os.environ.get("INVENTORY_DIR"),
        required=os.environ.get("INVENTORY_DIR") is None,
        help=(
            "Inventory directory base (default: $INVENTORY_DIR). "
            "When the included apps declare more than one matrix-deploy "
            "variant, sibling folders `<dir>-0`, `<dir>-1`, ... are "
            "created; otherwise the directory is used as-is."
        ),
    )

    g = p.add_mutually_exclusive_group(required=True)
    g.add_argument(
        "--apps",
        help="One or more application ids (will include run_after deps automatically).",
    )
    g.add_argument(
        "--include",
        help="Comma-separated list of application ids to include (no deps resolution).",
    )

    p.add_argument(
        "--threshold-gib",
        type=int,
        default=100,
        help="Free-space threshold (GiB) below which STORAGE_CONSTRAINED is enabled (default: 100).",
    )
    p.add_argument(
        "--force-storage-constrained",
        choices=["true", "false"],
        default=None,
        help="Override storage detection explicitly.",
    )
    p.add_argument(
        "--vars",
        default=None,
        help="JSON object merged into inventory variables (overrides win).",
    )
    p.add_argument(
        "--variant",
        type=int,
        default=_env_variant(),
        help=(
            "Pin the matrix init to a single round (zero-based index). "
            "Useful when only one variant of a multi-variant app needs to "
            "be (re-)materialised. Defaults to the VARIANT environment "
            "variable when set, otherwise full-matrix mode."
        ),
    )
    p.set_defaults(_handler=handler)


def _env_variant() -> int | None:
    raw = os.environ.get("VARIANT", "").strip()
    if not raw:
        return None
    try:
        return int(raw)
    except ValueError:
        raise SystemExit(
            f"VARIANT environment variable must be an integer, got {raw!r}"
        )


def handler(args: argparse.Namespace) -> int:
    compose = make_compose(distro=args.distro)

    if args.apps:
        include = resolve_deploy_ids_for_apps(compose, args.apps)
    else:
        include = [x.strip() for x in (args.include or "").split(",") if x.strip()]

    if not include:
        raise SystemExit("Resolved include list is empty")

    extra_vars: Dict[str, Any] | None = None
    if args.vars is not None:
        try:
            parsed = json.loads(args.vars)
        except Exception as exc:
            raise SystemExit(f"--vars must be valid JSON: {exc}")
        if not isinstance(parsed, dict):
            raise SystemExit("--vars must be a JSON object")
        extra_vars = parsed

    if args.force_storage_constrained is not None:
        storage_constrained = args.force_storage_constrained == "true"
    else:
        storage_constrained = detect_storage_constrained(
            compose, threshold_gib=int(args.threshold_gib)
        )

    plan = plan_dev_inventory_matrix(
        roles_dir=str(compose.repo_root / "roles"),
        include=include,
        base_inventory_dir=str(args.inventory_dir),
    )
    try:
        plan = filter_plan_to_variant(plan, args.variant)
    except ValueError as exc:
        raise SystemExit(f"--variant: {exc}")

    runtime = os.environ.get("RUNTIME") or detect_runtime()
    services_disabled = os.environ.get("SERVICES_DISABLED", "")
    for _round_index, inv_dir, round_variants in plan:
        spec = DevInventorySpec(
            inventory_dir=inv_dir,
            include=tuple(include),
            storage_constrained=storage_constrained,
            runtime=runtime,
            extra_vars=extra_vars,
            services_disabled=services_disabled,
            active_variants=round_variants,
        )
        build_dev_inventory(compose, spec)

    if len(plan) == 1:
        _, inv_dir, round_variants = plan[0]
        non_zero = {a: i for a, i in round_variants.items() if i}
        suffix = f" variants={non_zero}" if non_zero else ""
        print(
            f">>> Inventory initialized at {inv_dir} "
            f"(include={','.join(include)} storage_constrained={storage_constrained}){suffix}"
        )
    else:
        print(
            f">>> Matrix inventory initialized in {len(plan)} folders "
            f"(include={','.join(include)} storage_constrained={storage_constrained}):"
        )
        for round_index, inv_dir, round_variants in plan:
            non_zero = {a: i for a, i in round_variants.items() if i}
            print(
                f"    [round {round_index}] {inv_dir}"
                + (f"  variants={non_zero}" if non_zero else "")
            )
    return 0