computer-playbook/utils/github/variant/bundles.py

"""Expand a discovered app-id list into CI deploy-matrix entries.

A role whose matrix-deploy declares more variants than a single runner should
iterate is split into bundles of consecutive variant indices, one runner per
bundle — e.g. a 5-variant role becomes ``0,1,2`` and ``3,4``. A new bundle opens
as soon as the current one would exceed ``INFINITO_VARIANT_BUNDLE_SIZE`` variants
(default 3) OR ``INFINITO_VARIANT_BUNDLE_MAX_STORAGE`` cumulative ``min_storage``
(default 350GB), so storage-heavy variants are not stacked onto one runner that
would then run too long. Roles that fit one runner stay a single entry but still
carry their full ``0,…,N-1`` variant CSV, so every job name shows the variants it
covers.

Each entry is ``{"apps": <id>, "variant": "<csv>", "variant_slug": "<dashed>",
"label": "<display name>"}``; ``variant`` is empty only for a role that
declares no variants (a plain base deploy). ``label`` is the role's display
name (``utils.roles.display``) with the variant CSV appended, and is what the
deploy jobs title themselves with. The ``variant`` slice is threaded
through to ``cli.administration.deploy.development`` via the ``variant``
environment variable (consumed by ``--variant``), so a runner only iterates the
rounds in its bundle. ``variant_slug`` is a comma-free copy for artifact/job
names (GitHub Actions expressions have no string-replace function).

In swarm mode (``INFINITO_DEPLOY_MODE=swarm``) the discovery already selects
per variant (``role#variant`` tokens from cli.meta.ci.query); each token maps
1:1 to a runner. A bare role name (no ``#``) still expands to one entry per
variant. Compose mode packs several variants per runner.

Swarm and compose entries also carry the tor axis (see :func:`assign_tor`):
``tor`` says whether the entry deploys behind the node onion, ``disable``
carries the provider tokens the deploy drill switches off, and the label opens
with the onion glyph or the clearnet globe so a job name states which of the
two it tested. Host mode has no onion, so its entries are always assigned the
``disabled`` axis regardless of ``INFINITO_TOR`` and wear neither glyph.
"""

from __future__ import annotations

import json
import os
import sys
from typing import TYPE_CHECKING

from humanfriendly import InvalidSize, parse_size

from utils import PROJECT_ROOT
from utils.cache.applications import get_variants
from utils.cache.yaml import load_yaml_any
from utils.github.variant.limits import DEFAULT_BUNDLE_SIZE, DEFAULT_MAX_STORAGE
from utils.roles.applications.services.registry import (
    build_service_registry_from_applications,
    load_applications_from_roles_dir,
)
from utils.roles.applications.services.resources import (
    aggregate,
    collect_role_resources,
)
from utils.roles.display import display_names
from utils.roles.mapping import ROLE_FILE_META_SERVICES
from utils.symbol_glossary import to_emoji

if TYPE_CHECKING:
    from collections.abc import Iterable, Mapping, Sequence
    from pathlib import Path
    from typing import Any

ROLES_DIR = PROJECT_ROOT / "roles"

TOR_MODES = ("auto", "enforced", "exclusive", "disabled")

TOR_DEPLOY_MODES = ("swarm", "compose")


def _variant_indices(variant_csv: str) -> list[int]:
    """Parse a matrix entry's comma-separated variant indices."""
    return [
        int(token)
        for token in variant_csv.split(",")
        if token.strip().lstrip("-").isdigit()
    ]


def _tor_flag(config: Mapping[str, Any]) -> Any:
    """The ``services.tor.enabled`` value of one config, or ``None`` if unset."""
    services = config.get("services") if isinstance(config, dict) else None
    tor = services.get("tor") if isinstance(services, dict) else None
    return tor.get("enabled") if isinstance(tor, dict) else None


def _base_tor_capable(app: str) -> bool:
    """Read the tor gate straight from the role's ``meta/services.yml``."""
    path = ROLES_DIR / app / ROLE_FILE_META_SERVICES
    if not path.exists():
        return True
    try:
        services = load_yaml_any(path) or {}
    except Exception:  # noqa: BLE001  malformed role meta must not break the matrix
        return True
    return _tor_flag({"services": services}) is not False


def tor_capable(
    app: str,
    variant_csv: str = "",
    variants_per_app: Mapping[str, Sequence[Mapping[str, Any]]] | None = None,
) -> bool:
    """Whether a matrix entry may be deployed behind the node onion.

    Args:
        app: application id.
        variant_csv: the entry's variant indices, comma separated; empty for a
            role that declares none.
        variants_per_app: rendered variant configs per app; ``None`` falls back
            to the role's base ``meta/services.yml``.

    Returns:
        ``False`` when every variant the entry covers pins
        ``services.tor.enabled`` to a literal false, ``True`` otherwise. A
        reactive Jinja flag counts as capable: it is the role saying "onion
        when the node is dark". Consulting the variants matters — roles pin the
        gate ``true`` in variant 0 and ``false`` in the rest, so reading only
        the base config claims an onion for two thirds of the rounds that
        deliberately run without one.
    """
    declared = (variants_per_app or {}).get(app) or []
    covered = [
        declared[index]
        for index in _variant_indices(variant_csv)
        if 0 <= index < len(declared)
    ]
    if not covered:
        return _base_tor_capable(app)
    return any(_tor_flag(config) is not False for config in covered)


def keep_tor_variants(
    entries: list[dict[str, str]],
    variants_per_app: Mapping[str, Sequence[Mapping[str, Any]]] | None = None,
) -> list[dict[str, str]]:
    """Reduce every entry to the variants that can run behind the node onion.

    Args:
        entries: matrix entries in emission order.
        variants_per_app: rendered variant configs per app.

    Returns:
        a new list where each entry covers only its tor-capable variants, and
        entries left with none are gone. This is what ``exclusive`` means: the
        run tests the onion and nothing else, so a bundle of three variants of
        which one can use tor shrinks to that one instead of dragging the other
        two along. Unrelated to ``services.tor.exclusive``, which drops a
        deployed app's clearnet vhost; this one only narrows what CI runs.
    """
    kept: list[dict[str, str]] = []
    for entry in entries:
        app = entry["apps"]
        indices = _variant_indices(entry.get("variant", ""))
        if not indices:
            if tor_capable(app, "", variants_per_app):
                kept.append(entry)
            continue
        live = [
            index for index in indices if tor_capable(app, str(index), variants_per_app)
        ]
        if not live:
            continue
        csv = ",".join(str(index) for index in live)
        kept.append({**entry, "variant": csv, "variant_slug": csv.replace(",", "-")})
    return kept


def assign_tor(
    entries: list[dict[str, str]],
    mode: str,
    run: int,
    variants_per_app: Mapping[str, Sequence[Mapping[str, Any]]] | None = None,
) -> None:
    """Decide the tor axis for every matrix entry, in place.

    Args:
        entries: matrix entries in emission order.
        mode: ``enforced`` gives tor to every capable entry, ``disabled`` to
            none, ``auto`` to every second capable one. The run number's
            parity flips which one starts, so two consecutive runs cover
            both states on every capable entry. ``exclusive`` behaves like
            ``enforced``; the incapable entries are gone before it runs,
            dropped by :func:`keep_tor_variants`.
        run: run number.
        variants_per_app: rendered variant configs per app, so an entry whose
            variants switch the tor gate off is never counted capable.

    Every entry gains ``tor`` (``"true"``/``"false"``) and ``disable``, the
    token list handed to the deploy drill; an entry without tor disables the
    provider so no dependency edge can pull it back into the closure.
    """
    capable = [
        index
        for index, entry in enumerate(entries)
        if tor_capable(entry["apps"], entry.get("variant", ""), variants_per_app)
    ]
    if mode in ("enforced", "exclusive"):
        with_tor = set(capable)
    elif mode == "disabled":
        with_tor = set()
    else:
        with_tor = {
            index for position, index in enumerate(capable) if position % 2 == run % 2
        }
    for index, entry in enumerate(entries):
        enabled = index in with_tor
        entry["tor"] = "true" if enabled else "false"
        entry["disable"] = "" if enabled else "tor"


def resolve_tor_mode(raw: str | None = None) -> str:
    """Tor axis mode from ``INFINITO_TOR``; unknown or empty means ``auto``.

    Args:
        raw: explicit value; ``None`` reads the environment.

    Returns:
        one of :data:`TOR_MODES`.
    """
    if raw is None:
        raw = os.environ.get("INFINITO_TOR")
    value = (raw or "").strip().lower()
    return value if value in TOR_MODES else "auto"


def resolve_run_number(raw: str | None = None) -> int:
    """Run number from ``RUN_NUMBER``; its parity flips the ``auto`` split."""
    if raw is None:
        raw = os.environ.get("RUN_NUMBER")
    try:
        return int((raw or "0").strip())
    except ValueError:
        return 0


def resolve_bundle_size(raw: str | None = None) -> int:
    if raw is None:
        raw = os.environ.get("INFINITO_VARIANT_BUNDLE_SIZE")
    value = (raw or "").strip()
    if not value:
        return DEFAULT_BUNDLE_SIZE
    try:
        size = int(value)
    except ValueError:
        raise ValueError(
            f"INFINITO_VARIANT_BUNDLE_SIZE must be an integer, got {value!r}"
        ) from None
    if size < 1:
        raise ValueError("INFINITO_VARIANT_BUNDLE_SIZE must be >= 1")
    return size


def resolve_max_storage(raw: str | None = None) -> int | None:
    if raw is None:
        raw = os.environ.get("INFINITO_VARIANT_BUNDLE_MAX_STORAGE")
    value = (raw or "").strip()
    if not value:
        value = DEFAULT_MAX_STORAGE
    if value.lower() in ("0", "off", "none"):
        return None
    try:
        return int(parse_size(value))
    except InvalidSize:
        raise ValueError(
            f"INFINITO_VARIANT_BUNDLE_MAX_STORAGE must be a size like {DEFAULT_MAX_STORAGE!r}, "
            f"got {value!r}"
        ) from None


def _pack_indices(
    indices: Sequence[int],
    bundle_size: int,
    storages: Sequence[int | None] | None = None,
    max_storage_bytes: int | None = None,
) -> list[list[int]]:
    """Greedily pack the given (absolute) variant indices into bundles, opening
    a new bundle as soon as the current one would exceed ``bundle_size``
    variants OR ``max_storage_bytes`` cumulative ``min_storage``; both counters
    reset per bundle. ``storages`` is indexed by the ABSOLUTE variant index, so
    a filtered index list still reads the right per-variant storage."""
    bundles: list[list[int]] = []
    current: list[int] = []
    current_storage = 0
    for index in indices:
        size = 0
        if storages is not None and index < len(storages):
            size = storages[index] or 0
        over_count = len(current) >= bundle_size
        over_storage = (
            max_storage_bytes is not None
            and bool(current)
            and current_storage + size > max_storage_bytes
        )
        if current and (over_count or over_storage):
            bundles.append(current)
            current = []
            current_storage = 0
        current.append(index)
        current_storage += size
    if current:
        bundles.append(current)
    return bundles


def bundle_indices(
    count: int,
    bundle_size: int,
    storages: Sequence[int | None] | None = None,
    max_storage_bytes: int | None = None,
) -> list[list[int]]:
    """Pack ``0..count-1`` into bundles. With no storage cap this matches
    ``chunk_indices``. A single variant that alone exceeds the storage cap still
    forms its own bundle."""
    return _pack_indices(list(range(count)), bundle_size, storages, max_storage_bytes)


def chunk_indices(count: int, bundle_size: int) -> list[list[int]]:
    return [
        list(range(start, min(start + bundle_size, count)))
        for start in range(0, count, bundle_size)
    ]


def variant_count(variants_per_app: Mapping[str, Sequence[Any]], app: str) -> int:
    return max(1, len(variants_per_app.get(app) or []))


def _entry(app: str, variant_csv: str) -> dict[str, str]:
    """``variant_slug`` is a comma-free copy of ``variant_csv`` for artifact/job
    names: GitHub Actions expressions have no ``replace()``, so the CSV cannot be
    sanitised in YAML."""
    return {
        "apps": app,
        "variant": variant_csv,
        "variant_slug": variant_csv.replace(",", "-"),
    }


def expand_apps(
    apps: Iterable[str],
    variants_per_app: Mapping[str, Sequence[Any]],
    bundle_size: int = DEFAULT_BUNDLE_SIZE,
    *,
    storages_per_app: Mapping[str, Sequence[int | None]] | None = None,
    max_storage_bytes: int | None = None,
    bundle_size_per_app: Mapping[str, int] | None = None,
) -> list[dict[str, str]]:
    entries: list[dict[str, str]] = []
    for app in apps:
        variant_list = variants_per_app.get(app) or []
        if not variant_list:
            entries.append({"apps": app, "variant": "", "variant_slug": ""})
            continue
        indices = list(range(len(variant_list)))
        storages = (storages_per_app or {}).get(app)
        app_bundle_size = (bundle_size_per_app or {}).get(app) or bundle_size
        entries.extend(
            _entry(app, ",".join(map(str, chunk)))
            for chunk in _pack_indices(
                indices, app_bundle_size, storages, max_storage_bytes
            )
        )
    return entries


def app_bundle_sizes(apps: Iterable[str]) -> dict[str, int]:
    """Per-app ``variant_bundle_size`` overrides from ``meta/tests.yml``
    (runtime-heavy roles cap variants per compose job below the global
    default); apps without the key are absent."""
    from utils.roles.meta_lookup import get_role_variant_bundle_size

    sizes: dict[str, int] = {}
    for app in apps:
        role_dir = ROLES_DIR / app
        if not role_dir.is_dir():
            continue
        size = get_role_variant_bundle_size(role_dir, role_name=app)
        if size is not None:
            sizes[app] = size
    return sizes


def app_variant_storages(
    apps: Iterable[str],
    variants_per_app: Mapping[str, Sequence[Any]],
    roles_dir: Path = ROLES_DIR,
) -> dict[str, list[int | None]]:
    """Per-variant total ``min_storage`` (bytes) for each app, via the shared
    resource collection (same machinery as the ``ressources`` CLI / budget
    lint). Unknown apps and variants without storage yield ``None`` entries."""
    applications = load_applications_from_roles_dir(roles_dir)
    registry = build_service_registry_from_applications(applications)
    result: dict[str, list[int | None]] = {}
    for app in apps:
        sizes: list[int | None] = []
        for variant_config in variants_per_app.get(app) or []:
            scoped = dict(applications)
            scoped[app] = variant_config or {}
            rows: list[dict[str, Any]] = []
            collect_role_resources(
                role_name=app,
                applications=scoped,
                service_registry=registry,
                visited=set(),
                rows=rows,
                warnings=[],
                dedup=True,
            )
            sizes.append(aggregate(rows).get("min_storage_bytes"))
        result[app] = sizes
    return result


def compose_bundle_counts(
    apps: Iterable[str],
    variants_per_app: Mapping[str, Sequence[Any]],
    *,
    roles_dir: Path = ROLES_DIR,
) -> dict[str, int]:
    """Per app, the number of compose CI bundles (jobs) its variants pack into.

    Uses the same bundle-size + cumulative-``min_storage`` packing as the
    compose deploy matrix (``expand_apps`` / ``bundle_indices``), so the
    ``complexity`` report and the matrix never diverge on the job count.
    """
    apps = list(apps)
    storages = app_variant_storages(apps, variants_per_app, roles_dir)
    bundle_size = resolve_bundle_size()
    max_storage = resolve_max_storage()
    overrides = app_bundle_sizes(apps)
    return {
        app: len(
            bundle_indices(
                variant_count(variants_per_app, app),
                overrides.get(app) or bundle_size,
                storages.get(app),
                max_storage,
            )
        )
        for app in apps
    }


def _deploy_mode() -> str:
    return (os.environ.get("INFINITO_DEPLOY_MODE") or "").strip().lower()


def _swarm_mode() -> bool:
    return _deploy_mode() == "swarm"


def main(argv: Sequence[str] | None = None) -> int:
    argv = list(sys.argv[1:] if argv is None else argv)
    raw = argv[0] if argv else sys.stdin.read()
    apps = json.loads(raw) if raw.strip() else []
    if not isinstance(apps, list):
        raise SystemExit(
            f"variant.bundles: expected a JSON array of app ids, got "
            f"{type(apps).__name__}"
        )
    variants_per_app = get_variants()
    if _swarm_mode():
        entries = []
        for app in apps:
            if "#" in app:
                entries.append(_entry(*app.split("#", 1)))
            else:
                entries.extend(expand_apps([app], variants_per_app, 1))
    else:
        entries = expand_apps(
            apps,
            variants_per_app,
            resolve_bundle_size(),
            storages_per_app=app_variant_storages(apps, variants_per_app),
            max_storage_bytes=resolve_max_storage(),
            bundle_size_per_app=app_bundle_sizes(apps),
        )
    on_axis = _deploy_mode() in TOR_DEPLOY_MODES
    tor_mode = resolve_tor_mode() if on_axis else "disabled"
    if tor_mode == "exclusive":
        entries = keep_tor_variants(entries, variants_per_app)
    assign_tor(entries, tor_mode, resolve_run_number(), variants_per_app)
    codec = display_names()
    for entry in entries:
        label = codec.encode(entry["apps"], entry["variant"])
        if on_axis:
            glyph = to_emoji("tor" if entry["tor"] == "true" else "clearnet")
            label = f"{glyph}{label}"
        entry["label"] = label
    print(json.dumps(entries))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())