devx/src/devx/molecule/distribute_molecule.py

#!/usr/bin/env python3
"""Distribute molecule (scenario, platform) pairs across N parallel runners.

Discovers all molecule scenarios under ansible/roles/*/molecule/ and
crosses them with the supported OS platform matrix, then splits the
resulting test pairs evenly across the requested number of runners.

Each pair is printed as ``scenario|platform_name|platform_image|platform_command``
so the CI workflow can set the appropriate environment variables.

Usage:
  python3 -m devx.molecule.distribute_molecule --runner-index 1 --max-runners 3
  # prints: default|ubuntu-2204|ubuntu:22.04| lifecycle|ubuntu-2204|ubuntu:22.04| ...
  python3 -m devx.molecule.distribute_molecule --list
  # prints all scenarios, one per line
  python3 -m devx.molecule.distribute_molecule --list-platforms
  # prints all platforms, one per line
"""

from __future__ import annotations

import tomllib
from dataclasses import dataclass
from pathlib import Path

import click

from devx.i18n import _
from devx.molecule.platforms import PLATFORMS, load_platforms

DEFAULT_MAX_RUNNERS = 3
MOLECULE_ROOT = Path("ansible/roles/gitea-runner/molecule")
DEFAULT_ROLES_ROOT = Path("ansible/roles")


@dataclass(frozen=True)
class TestPair:
    """A (scenario, platform) combination to test."""

    scenario: str
    platform: dict[str, str]

    def encode(self) -> str:
        """Serialize to a pipe-delimited string for CI consumption."""
        cmd = self.platform["command"].replace(" ", "__SPACE__")
        return f"{self.scenario}|{self.platform['name']}|{self.platform['image']}|{cmd}"

    @staticmethod
    def decode(encoded: str) -> TestPair:
        """Deserialize from a pipe-delimited string."""
        parts = encoded.split("|")
        return TestPair(
            scenario=parts[0],
            platform={"name": parts[1], "image": parts[2], "command": parts[3].replace("__SPACE__", " ")},
        )


@dataclass(frozen=True)
class MultiRoleTestPair:
    """A (role, scenario, platform) combination for multi-role projects."""

    role: str
    scenario: str
    platform: dict[str, str]

    def encode(self) -> str:
        """Serialize to a pipe-delimited string: ``role|scenario|platform_name|image|command``."""
        cmd = self.platform["command"].replace(" ", "__SPACE__")
        return f"{self.role}|{self.scenario}|{self.platform['name']}|{self.platform['image']}|{cmd}"

    @staticmethod
    def decode(encoded: str) -> MultiRoleTestPair:
        """Deserialize from a pipe-delimited string."""
        parts = encoded.split("|")
        return MultiRoleTestPair(
            role=parts[0],
            scenario=parts[1],
            platform={"name": parts[2], "image": parts[3], "command": parts[4].replace("__SPACE__", " ")},
        )


def discover_scenarios(root: Path | None = None) -> list[str]:
    """Return sorted list of molecule scenario directory names."""
    if root is None:
        root = MOLECULE_ROOT
    if not root.is_dir():
        raise click.ClickException(_("Molecule directory not found: {path}", path=str(root)))
    scenarios = [d.name for d in root.iterdir() if d.is_dir() and not d.name.startswith("_") and d.name != "common"]
    return sorted(scenarios)


def discover_multi_role_scenarios(roles_root: Path | None = None) -> list[tuple[str, str]]:
    """Discover (role, scenario) pairs across all roles under *roles_root*.

    Scans ``roles_root/*/molecule/*/`` for scenario directories, skipping
    ``common`` and directories starting with ``_``. Returns a sorted list of
    ``(role_name, scenario_name)`` tuples.
    """
    if roles_root is None:
        roles_root = DEFAULT_ROLES_ROOT
    if not roles_root.is_dir():
        raise click.ClickException(_("Roles directory not found: {path}", path=str(roles_root)))
    pairs: list[tuple[str, str]] = []
    for role_dir in sorted(roles_root.iterdir()):
        if not role_dir.is_dir():
            continue
        mol_dir = role_dir / "molecule"
        if not mol_dir.is_dir():
            continue
        for scenario_dir in mol_dir.iterdir():
            if not scenario_dir.is_dir():
                continue
            if scenario_dir.name.startswith("_") or scenario_dir.name == "common":
                continue
            pairs.append((role_dir.name, scenario_dir.name))
    return pairs


def build_pairs(scenarios: list[str], platforms: list[dict[str, str]] | None = None) -> list[TestPair]:
    """Build the full cross-product of scenarios and platforms."""
    if platforms is None:
        platforms = PLATFORMS
    return [TestPair(s, p) for s in scenarios for p in platforms]


def build_multi_role_pairs(
    role_scenarios: list[tuple[str, str]],
    platforms: list[dict[str, str]] | None = None,
) -> list[MultiRoleTestPair]:
    """Build the full cross-product of (role, scenario) pairs and platforms."""
    if platforms is None:
        platforms = PLATFORMS
    return [MultiRoleTestPair(r, s, p) for r, s in role_scenarios for p in platforms]


# --- Molecule weight configuration ---
#
# Weights are loaded from ``[tool.devx.molecule.weights]`` in
# ``pyproject.toml``.  Each project (infra, grm, …) contributes its own
# weights calibrated from actual CI execution times.
#
# Two key formats are supported:
#   - ``"scenario" = weight`` — applies to any role with that scenario name
#   - ``"role/scenario" = weight`` — role-specific (takes priority)
#
# Example pyproject.toml::
#
#   [tool.devx.molecule.weights]
#   "nextcloud" = 15
#   "app_container/customer-apps" = 11
#   "restore/default" = 11
#   "default" = 3
#
# If no configuration is found, a generic default weight is used for all
# scenarios (producing a round-robin distribution).

_DEFAULT_SCENARIO_WEIGHT = 3


def _load_molecule_weights(pyproject_path: str = "pyproject.toml") -> tuple[dict[str, int], dict[tuple[str, str], int]]:
    """Load molecule weights from ``[tool.devx.molecule.weights]`` in pyproject.toml.

    Returns a tuple of ``(scenario_weights, role_scenario_weights)``:
    - ``scenario_weights``: maps scenario name → weight (applies to any role)
    - ``role_scenario_weights``: maps (role, scenario) → weight (role-specific)
    """
    path = Path(pyproject_path)
    if not path.exists():
        return {}, {}
    try:
        with open(path, "rb") as f:  # noqa: PTH123
            data = tomllib.load(f)
    except (tomllib.TOMLDecodeError, OSError):
        return {}, {}

    weights_raw = data.get("tool", {}).get("devx", {}).get("molecule", {}).get("weights", {})
    if not isinstance(weights_raw, dict):
        return {}, {}

    scenario_weights: dict[str, int] = {}
    role_scenario_weights: dict[tuple[str, str], int] = {}

    for key, value in weights_raw.items():
        if not isinstance(value, int):
            continue
        if "/" in key:
            role, scenario = key.split("/", 1)
            role_scenario_weights[(role.lower(), scenario.lower())] = value
        else:
            scenario_weights[key.lower()] = value

    return scenario_weights, role_scenario_weights


# Load weights once at import time (like devx.config and classify_changes)
_SCENARIO_WEIGHTS, _ROLE_SCENARIO_WEIGHTS = _load_molecule_weights()


def _scenario_weight(scenario: str, role: str | None = None) -> int:
    """Estimate a weight for a scenario based on its name and optionally its role.

    Role-specific weights (``"role/scenario"``) take priority over
    scenario-name-only weights (``"scenario"``).  Falls back to the
    default weight if no configuration matches.
    """
    s = scenario.lower()
    if role is not None:
        r = role.lower()
        key = (r, s)
        if key in _ROLE_SCENARIO_WEIGHTS:
            return _ROLE_SCENARIO_WEIGHTS[key]
    for key, weight in _SCENARIO_WEIGHTS.items():
        if key in s:
            return weight
    return _DEFAULT_SCENARIO_WEIGHT


def _lpt_distribute[T](items: list[T], weights: list[int], max_runners: int) -> list[list[T]]:
    """Distribute *items* across *max_runners* using LPT (Longest Processing Time first).

    Sorts items by weight (descending), then assigns each to the runner
    with the least total weight. This produces a more balanced distribution
    than naive round-robin when items have varying costs.
    """
    groups: list[list[T]] = [[] for _ in range(max_runners)]
    loads = [0] * max_runners
    # Sort by weight descending, preserving original order for ties
    indexed = sorted(enumerate(items), key=lambda x: (-weights[x[0]], x[0]))
    for orig_idx, item in indexed:
        # Find the runner with the minimum load
        min_runner = min(range(max_runners), key=lambda r: loads[r])
        groups[min_runner].append(item)
        loads[min_runner] += weights[orig_idx]
    return groups


def distribute_multi_role(pairs: list[MultiRoleTestPair], max_runners: int) -> list[list[MultiRoleTestPair]]:
    """Split *pairs* into *max_runners* balanced groups using LPT scheduling.

    Each pair is weighted by role+scenario heuristics (e.g. ``nextcloud`` is
    heavier than ``simple-app``). Pairs are sorted by weight descending and
    assigned to the runner with the least total weight.
    """
    weights = [_scenario_weight(p.scenario, p.role) for p in pairs]
    return _lpt_distribute(pairs, weights, max_runners)


def multi_role_pairs_for_runner(
    pairs: list[MultiRoleTestPair], runner_index: int, max_runners: int
) -> list[MultiRoleTestPair]:
    """Return the subset of multi-role pairs assigned to *runner_index* (0-based)."""
    groups = distribute_multi_role(pairs, max_runners)
    if runner_index < 0 or runner_index >= len(groups):
        raise click.ClickException(
            _("Runner index {index} out of range (0..{max})", index=runner_index, max=max_runners - 1)
        )
    return groups[runner_index]


def distribute(pairs: list[TestPair], max_runners: int) -> list[list[TestPair]]:
    """Split *pairs* into *max_runners* balanced groups using LPT scheduling.

    Each pair is weighted by scenario name heuristics (e.g. ``nextcloud`` is
    heavier than ``binary``). Pairs are sorted by weight descending and
    assigned to the runner with the least total weight.
    """
    weights = [_scenario_weight(p.scenario) for p in pairs]
    return _lpt_distribute(pairs, weights, max_runners)


def pairs_for_runner(pairs: list[TestPair], runner_index: int, max_runners: int) -> list[TestPair]:
    """Return the subset of pairs assigned to *runner_index*."""
    groups = distribute(pairs, max_runners)
    if runner_index < 0 or runner_index >= len(groups):
        raise click.ClickException(
            _(
                "Runner index {index} out of range (0..{max})",
                index=runner_index,
                max=max_runners - 1,
            )
        )
    return groups[runner_index]


def _write_github_env(key: str, value: str) -> None:
    """Append a key=value line to the $GITHUB_ENV file."""
    import os

    gh_env = os.environ.get("GITHUB_ENV")
    if not gh_env:
        raise click.ClickException("GITHUB_ENV environment variable is not set")
    with open(gh_env, "a") as f:  # noqa: PTH123
        f.write(f"{key}={value}\n")


@click.command()
@click.option(
    "--runner-index",
    type=int,
    default=None,
    help="One-based runner index (Gitea Actions renders 0 as empty). "
    "Converted to zero-based internally. If omitted, prints all groups.",
)
@click.option(
    "--max-runners",
    type=int,
    default=DEFAULT_MAX_RUNNERS,
    show_default=True,
    help="Total number of parallel runners.",
)
@click.option(
    "--list",
    "list_all",
    is_flag=True,
    help="List all discovered scenarios, one per line.",
)
@click.option(
    "--list-platforms",
    "list_platforms",
    is_flag=True,
    help="List all supported platforms, one per line.",
)
@click.option(
    "--github-env",
    "github_env",
    is_flag=True,
    default=False,
    help="Write TEST_PAIRS and SKIP to $GITHUB_ENV (for CI workflow steps).",
)
@click.option(
    "--skip-if-excess",
    is_flag=True,
    default=False,
    help="With --github-env: write SKIP=true when runner-index exceeds max-runners.",
)
@click.option(
    "--molecule-root",
    type=click.Path(exists=True, file_okay=False, path_type=Path),
    default=None,
    help="Custom molecule directory (single-role mode). Default: ansible/roles/gitea-runner/molecule.",
)
@click.option(
    "--roles-root",
    type=click.Path(exists=True, file_okay=False, path_type=Path),
    default=None,
    help="Roles directory for multi-role discovery (scans */molecule/*/). "
    "Use this for projects with multiple Ansible roles. Default: disabled (single-role mode).",
)
@click.option(
    "--platforms-file",
    type=click.Path(exists=True, file_okay=True, path_type=Path),
    default=None,
    help="JSON file with custom platform list (each entry: name, image, command). "
    "Overrides the default platform matrix. Useful for projects with custom test images.",
)
def cli(
    runner_index: int | None,
    max_runners: int,
    list_all: bool,
    list_platforms: bool,
    github_env: bool,
    skip_if_excess: bool,
    molecule_root: Path | None,
    roles_root: Path | None,
    platforms_file: Path | None,
) -> None:
    platforms = load_platforms(platforms_file)
    # Multi-role mode: discover (role, scenario) pairs across all roles
    if roles_root is not None:
        role_scenarios = discover_multi_role_scenarios(roles_root)
        if list_all:
            for role, scenario in role_scenarios:
                click.echo(f"{role}|{scenario}")
            return
        if list_platforms:
            for p in platforms:
                click.echo(f"{p['name']}|{p['image']}|{p['command']}")
            return
        pairs_mr = build_multi_role_pairs(role_scenarios, platforms)
        if runner_index is None:
            groups = distribute_multi_role(pairs_mr, max_runners)
            for i, group in enumerate(groups):
                labels = " ".join(p.encode() for p in group) if group else "(none)"
                click.echo(f"Runner {i}: {labels}")
            return
        if skip_if_excess and github_env and runner_index > max_runners:
            click.echo(f"Skipping — runner index {runner_index} > max runners {max_runners}")
            _write_github_env("TEST_PAIRS", "")
            _write_github_env("SKIP", "true")
            return
        if runner_index < 1:
            raise click.ClickException(f"Runner index {runner_index} is out of range (must be >= 1)")
        zero_based = runner_index - 1
        assigned = multi_role_pairs_for_runner(pairs_mr, zero_based, max_runners)
        encoded = " ".join(p.encode() for p in assigned)
        if github_env:
            _write_github_env("TEST_PAIRS", encoded)
            _write_github_env("SKIP", "false")
            click.echo(f"Assigned pairs: {encoded}")
            return
        click.echo(encoded)
        return

    # Single-role mode (default or --molecule-root)
    scenarios = discover_scenarios(molecule_root)
    if list_all:
        for s in scenarios:
            click.echo(s)
        return
    if list_platforms:
        for p in platforms:
            click.echo(f"{p['name']}|{p['image']}|{p['command']}")
        return
    pairs = build_pairs(scenarios, platforms)
    if runner_index is None:
        groups = distribute(pairs, max_runners)
        for i, group in enumerate(groups):
            labels = " ".join(p.encode() for p in group) if group else "(none)"
            click.echo(f"Runner {i}: {labels}")
        return

    # Skip if runner index exceeds available runners (CI static matrix has 3 slots)
    if skip_if_excess and github_env and runner_index > max_runners:
        click.echo(f"Skipping — runner index {runner_index} > max runners {max_runners}")
        _write_github_env("TEST_PAIRS", "")
        _write_github_env("SKIP", "true")
        return

    # Validate runner index is in range
    if runner_index < 1:
        raise click.ClickException(f"Runner index {runner_index} is out of range (must be >= 1)")

    # Convert 1-based CLI index to 0-based internal index
    zero_based = runner_index - 1
    assigned = pairs_for_runner(pairs, zero_based, max_runners)
    encoded = " ".join(p.encode() for p in assigned)

    if github_env:
        _write_github_env("TEST_PAIRS", encoded)
        _write_github_env("SKIP", "false")
        click.echo(f"Assigned pairs: {encoded}")
        return

    click.echo(encoded)


if __name__ == "__main__":  # pragma: no cover
    cli()