computer-playbook/utils/roles/applications/services/registry.py

from __future__ import annotations

import functools
from pathlib import Path
from typing import TYPE_CHECKING, Any

from utils.cache.yaml import load_yaml_any
from utils.roles.entity.name import get_entity_name
from utils.roles.mapping import ROLE_FILE_VARS_MAIN
from utils.roles.meta_lookup import get_role_run_after
from utils.roles.validation.invokable import types_from_group_names

if TYPE_CHECKING:
    from collections.abc import Callable, Iterable


class ServiceRegistryError(ValueError):
    """Raised when role-local service discovery is invalid."""


def _as_mapping(value: Any) -> dict[str, Any]:
    return value if isinstance(value, dict) else {}


def _normalized_name(value: Any) -> str:
    if not isinstance(value, str):
        return ""
    return value.strip()


GROUP_GATED_CONSTANTS: tuple[str, ...] = ("LITELLM_USABLE",)
"""Group_vars constants whose own definition carries an ``in group_names``
gate, so a flag referencing one is dynamic even though its own text is not.

A constant hides the role id the substring rule looks for, so every SPOT of
this kind MUST be listed here or the roles referencing it silently drop their
provider from the static closure.
"""


def is_explicit_truth(value: Any) -> bool:
    """A services.<key> flag (``enabled`` / ``shared``) MAY be true if it
    is the literal Python ``True``, any string containing the substring
    ``in group_names``, or a string referencing a `GROUP_GATED_CONSTANTS`
    entry. Negated and compound predicates
    (``"{{ 'X' not in group_names }}"``) match on purpose: they are true
    on some node shape, and the static closure is the co-deploy superset.
    Mirrors the substring test at
    ``tests/integration/roles/meta/services/test_dynamic_flags.py:194``.

    Static-analysis callers (``ServicesResolver``,
    ``resolve_service_dependency_roles_from_config``, lint tests) MUST
    use this helper instead of comparing against ``True`` directly so
    that the dynamic-flag form keeps resolving deps at tooling time.
    The runtime path goes through ansible templating before reaching
    these checks, where the Jinja has already been rendered to a real
    boolean, so accepting the unrendered form is purely additive.
    """
    if value is True:
        return True
    if not isinstance(value, str):
        return False
    return "in group_names" in value or any(
        constant in value for constant in GROUP_GATED_CONSTANTS
    )


def detect_service_channel(role_name: str) -> str:
    return "frontend" if role_name.startswith(("web-app-", "web-svc-")) else "backend"


def detect_deploy_type(role_name: str) -> str:
    detected = types_from_group_names([role_name])
    if "server" in detected:
        return "server"
    if "workstation" in detected:
        return "workstation"
    if "universal" in detected:
        return "universal"
    return "server" if role_name.startswith(("web-app-", "web-svc-")) else "universal"


def detect_service_bucket(role_name: str) -> str:
    deploy_type = detect_deploy_type(role_name)
    if deploy_type == "universal":
        return "universal"
    if deploy_type == "workstation":
        return "workstation"
    if role_name.startswith("web-svc-"):
        return "web-svc"
    if role_name.startswith("web-app-"):
        return "web-app"
    return deploy_type


def read_yaml_file(path: Path) -> dict[str, Any]:
    if not path.is_file():
        return {}
    return _as_mapping(load_yaml_any(str(path), default_if_missing={}) or {})


def load_applications_from_roles_dir(roles_dir: Path) -> dict[str, dict[str, Any]]:
    applications: dict[str, dict[str, Any]] = {}
    for role_dir in sorted(p for p in roles_dir.iterdir() if p.is_dir()):
        vars_file = role_dir / ROLE_FILE_VARS_MAIN
        if not vars_file.is_file():
            continue
        vars_data = read_yaml_file(vars_file)
        application_id = _normalized_name(vars_data.get("application_id"))
        if not application_id:
            continue
        meta_dir = role_dir / "meta"
        config: dict[str, Any] = {}
        services_data = read_yaml_file(meta_dir / "services.yml")
        if services_data:
            config["services"] = services_data
        for topic in ("server", "rbac", "volumes"):
            topic_data = read_yaml_file(meta_dir / f"{topic}.yml")
            if topic_data:
                config[topic] = topic_data
        applications[application_id] = config
    return applications


def discover_role_services(
    role_name: str,
    config: dict[str, Any],
) -> dict[str, dict[str, Any]]:
    services = _as_mapping(config.get("services"))
    entity_name = get_entity_name(role_name)
    primary_entry = _as_mapping(services.get(entity_name))
    alias_entries = {
        key: _as_mapping(entry)
        for key, entry in services.items()
        if isinstance(entry, dict)
        and _normalized_name(_as_mapping(entry).get("canonical")) == entity_name
    }
    provides = _normalized_name(primary_entry.get("provides"))
    if provides == entity_name:
        provides = ""

    is_provider = bool(primary_entry) and (
        bool(primary_entry.get("shared"))
        or "provides" in primary_entry
        or alias_entries
    )
    if not is_provider:
        return {}

    primary_id = provides or entity_name
    raw_covers = primary_entry.get("covers")
    covers: list[str] = (
        [_normalized_name(item) for item in raw_covers if isinstance(item, str)]
        if isinstance(raw_covers, list)
        else []
    )
    covers = [c for c in covers if c]
    base_entry = {
        "role": role_name,
        "entity_name": entity_name,
        "source_key": entity_name,
        "deploy_type": detect_deploy_type(role_name),
        "bucket": detect_service_bucket(role_name),
        "service_type": detect_service_channel(role_name),
        "shared": bool(primary_entry.get("shared", False)),
        "enabled": bool(primary_entry.get("enabled", False)),
        "covers": covers,
    }
    app_networks = _as_mapping(config.get("networks"))
    overlay = app_networks.get("overlay")
    if isinstance(overlay, dict):
        base_entry["overlay"] = overlay
    if provides:
        base_entry["provides"] = provides

    discovered: dict[str, dict[str, Any]] = {primary_id: base_entry}
    for alias_key, alias_entry in sorted(alias_entries.items()):
        discovered[alias_key] = {
            **base_entry,
            "source_key": alias_key,
            "canonical": primary_id,
            "shared": bool(alias_entry.get("shared", False)),
            "enabled": bool(alias_entry.get("enabled", False)),
        }

    return discovered


def build_service_registry_from_applications(
    applications: dict[str, Any],
) -> dict[str, dict[str, Any]]:
    registry: dict[str, dict[str, Any]] = {}
    for role_name, config in sorted(applications.items()):
        role_services = discover_role_services(role_name, _as_mapping(config))
        for service_key, entry in role_services.items():
            existing = registry.get(service_key)
            if existing and existing.get("role") != entry.get("role"):
                raise ServiceRegistryError(
                    f"Duplicate service key '{service_key}' is declared by both "
                    f"'{existing.get('role')}' and '{entry.get('role')}'."
                )
            registry[service_key] = entry
    return registry


@functools.cache
def _service_registry_by_roles_dir(roles_dir_key: str) -> dict[str, dict[str, Any]]:
    return build_service_registry_from_applications(
        load_applications_from_roles_dir(Path(roles_dir_key))
    )


def build_service_registry_from_roles_dir(
    roles_dir: Path,
) -> dict[str, dict[str, Any]]:
    return _service_registry_by_roles_dir(str(Path(roles_dir).resolve()))


def expand_service_tokens(tokens: list[str], roles_dir: Path) -> list[str]:
    """Normalize service reference tokens so a caller may name either the
    service/entity key (e.g. ``email``, ``prometheus``) or the provider
    application id (e.g. ``web-app-mailu``, ``web-app-prometheus``).

    An application id is expanded to the service key(s) that role provides;
    every other token is kept verbatim. Order-preserving, de-duplicated.
    Used wherever the operator supplies service references (``disable``,
    include/prune closures) so both spellings resolve identically.
    """
    if not tokens:
        return []
    try:
        registry = build_service_registry_from_roles_dir(roles_dir)
    except Exception:  # noqa: BLE001 - normalization must never crash a caller; fall back to the raw tokens
        return list(dict.fromkeys(tokens))

    role_to_keys: dict[str, list[str]] = {}
    for key, entry in registry.items():
        role = (entry or {}).get("role")
        if isinstance(role, str) and role:
            role_to_keys.setdefault(role, []).append(key)

    out: list[str] = []
    for token in tokens:
        if token in registry:
            out.append(token)
        elif token in role_to_keys:
            out.extend(role_to_keys[token])
        else:
            out.append(token)
    return list(dict.fromkeys(out))


def build_role_to_primary_service_key(
    service_registry: dict[str, dict[str, Any]],
) -> dict[str, str]:
    result: dict[str, str] = {}
    for service_key, entry in service_registry.items():
        if "canonical" in entry:
            continue
        role_name = _normalized_name(entry.get("role"))
        if role_name:
            result[role_name] = service_key
    return result


def build_role_to_covered_keys(
    service_registry: dict[str, dict[str, Any]],
) -> dict[str, list[str]]:
    """For each role, return ``[primary_key, *covers]`` — the set of
    service keys whose presence in a consumer's ``services`` map counts
    as an explicit declaration of a dep on that role. ``covers`` comes
    from the primary entity's ``covers:`` list (e.g. ``web-app-keycloak``
    covers ``oauth2`` because the per-app oauth2-proxy only makes sense
    when the IdP is on the host)."""
    result: dict[str, list[str]] = {}
    for service_key, entry in service_registry.items():
        if "canonical" in entry:
            continue
        role_name = _normalized_name(entry.get("role"))
        if not role_name:
            continue
        keys = [service_key]
        for covered in entry.get("covers", []) or []:
            covered_name = _normalized_name(covered)
            if covered_name and covered_name not in keys:
                keys.append(covered_name)
        result[role_name] = keys
    return result


def build_covered_key_to_role(
    service_registry: dict[str, dict[str, Any]],
) -> dict[str, str]:
    """Inverse of :func:`build_role_to_covered_keys`: for each covered
    service key (NOT including the primary), the role that covers it.
    Use :func:`build_role_to_primary_service_key` for the primary
    direction; this helper is intentionally limited to the
    ``covers:`` extension so callers can distinguish "this role IS the
    provider" from "this role covers this dep transitively"."""
    result: dict[str, str] = {}
    for entry in service_registry.values():
        if "canonical" in entry:
            continue
        role_name = _normalized_name(entry.get("role"))
        if not role_name:
            continue
        for covered in entry.get("covers", []) or []:
            covered_name = _normalized_name(covered)
            if covered_name and covered_name not in result:
                result[covered_name] = role_name
    return result


def canonical_service_key(
    service_registry: dict[str, dict[str, Any]],
    service_key: str,
) -> str:
    entry = _as_mapping(service_registry.get(service_key))
    return _normalized_name(entry.get("canonical")) or service_key


def equivalent_service_keys(
    service_registry: dict[str, dict[str, Any]],
    service_key: str,
) -> list[str]:
    primary = canonical_service_key(service_registry, service_key)
    keys = [
        key
        for key, entry in service_registry.items()
        if canonical_service_key(service_registry, key) == primary
    ]
    return sorted(keys)


def resolve_service_dependency_roles_from_config(
    config: dict[str, Any],
    service_registry: dict[str, dict[str, Any]],
) -> list[str]:
    services = _as_mapping(config.get("services"))
    includes: list[str] = []
    for service_key, raw_service_conf in services.items():
        service_conf = _as_mapping(raw_service_conf)
        if not (
            is_explicit_truth(service_conf.get("enabled"))
            and is_explicit_truth(service_conf.get("shared"))
        ):
            continue

        entry = _as_mapping(service_registry.get(service_key))
        role_name = _normalized_name(entry.get("role"))
        if role_name:
            includes.append(role_name)

    seen = set()
    ordered: list[str] = []
    for role_name in includes:
        if role_name not in seen:
            ordered.append(role_name)
            seen.add(role_name)
    return ordered


def load_run_after_from_roles_dir(roles_dir: Path, role_name: str) -> list[str]:
    try:
        result = get_role_run_after(roles_dir / role_name, role_name=role_name)
    except Exception as exc:
        raise ServiceRegistryError(
            f"Invalid run_after in roles/{role_name}/meta/services.yml: {exc}"
        ) from exc
    return result


_BUCKET_ORDER = {
    "universal": 0,
    "workstation": 1,
    "server": 2,
    "web-svc": 3,
    "web-app": 4,
}


def bucket_order_key(role_name: str, bucket: str | None) -> tuple[int, str]:
    """Tie-break key for a ready role: coarse bucket first, then name.

    The bucket is a name-prefix-derived loader phase; it only breaks ties
    among roles with no pending run_after prerequisite, so run_after stays
    bucket-agnostic (an edge orders across any bucket).
    """
    return (_BUCKET_ORDER.get(bucket, len(_BUCKET_ORDER)), role_name)


def run_after_topological_order(
    nodes: Iterable[str],
    run_after_of: Callable[[str], Iterable[str]],
    tiebreak_of: Callable[[str], Any],
) -> list[str]:
    """Order ``nodes`` so every run_after prerequisite precedes its dependents.

    run_after edges order across any category; ``tiebreak_of`` only decides
    among nodes that are simultaneously ready (no pending prerequisite), so a
    coarse phase can stay the default for edge-less roles without constraining
    run_after itself.

    Args:
        nodes: role names to order (deduplicated internally).
        run_after_of: role name -> its run_after role names.
        tiebreak_of: role name -> sort key for ready-set tie-breaking.

    Raises:
        ServiceRegistryError: on a run_after cycle.
    """
    node_list = list(dict.fromkeys(nodes))
    node_set = set(node_list)
    graph: dict[str, list[str]] = {node: [] for node in node_list}
    indegree: dict[str, int] = dict.fromkeys(node_list, 0)

    for node in node_list:
        for dep in run_after_of(node):
            if dep not in node_set:
                continue
            graph[dep].append(node)
            indegree[node] += 1

    ready = [node for node, count in indegree.items() if count == 0]
    ordered: list[str] = []
    while ready:
        ready.sort(key=tiebreak_of)
        node = ready.pop(0)
        ordered.append(node)
        for dependent in sorted(graph[node]):
            indegree[dependent] -= 1
            if indegree[dependent] == 0:
                ready.append(dependent)

    if len(ordered) != len(node_set):
        raise ServiceRegistryError("Circular run_after dependency detected.")

    return ordered


def ordered_primary_service_entries(
    service_registry: dict[str, dict[str, Any]],
    roles_dir: Path,
) -> list[dict[str, Any]]:
    primary_entries = {
        entry["role"]: {"id": service_key, **entry}
        for service_key, entry in service_registry.items()
        if "canonical" not in entry
    }

    order = run_after_topological_order(
        primary_entries,
        lambda role: load_run_after_from_roles_dir(roles_dir, role),
        lambda role: bucket_order_key(role, primary_entries[role].get("bucket")),
    )
    return [primary_entries[role] for role in order]