#!/usr/bin/env python3
"""Distribute molecule scenarios across N parallel runners.

Discovers all molecule scenarios under ansible/roles/*/molecule/ and
splits them evenly across the requested number of runners.

Usage:
  python3 scripts/distribute_molecule.py --runner-index 0 --max-runners 3
  # prints: default deregister
  python3 scripts/distribute_molecule.py --list
  # prints all scenarios, one per line
"""

from __future__ import annotations

import sys
from pathlib import Path

import click

from gitea_runner_manager.i18n import _

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


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 distribute(scenarios: list[str], max_runners: int) -> list[list[str]]:
    """Split *scenarios* into *max_runners* balanced groups (round-robin)."""
    groups: list[list[str]] = [[] for _ in range(max_runners)]
    for i, scenario in enumerate(scenarios):
        groups[i % max_runners].append(scenario)
    return groups


def scenarios_for_runner(
    scenarios: list[str], runner_index: int, max_runners: int
) -> list[str]:
    """Return the subset of scenarios assigned to *runner_index*."""
    groups = distribute(scenarios, 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]


@click.command()
@click.option(
    "--runner-index",
    type=int,
    default=None,
    help="Zero-based runner index. 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.",
)
def cli(runner_index: int | None, max_runners: int, list_all: bool) -> None:
    scenarios = discover_scenarios()
    if list_all:
        for s in scenarios:
            click.echo(s)
        return
    if runner_index is None:
        groups = distribute(scenarios, max_runners)
        for i, group in enumerate(groups):
            click.echo(f"Runner {i}: {' '.join(group) if group else '(none)'}")
        return
    assigned = scenarios_for_runner(scenarios, runner_index, max_runners)
    click.echo(" ".join(assigned))


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