Kairo-Jinkendo/backend/steering/methods/registry.py

"""Method Registry — built-in steering methods — AP2.3e / AP2.4."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Literal

from steering.elements.registry import validate_steering_element_keys
from steering.graph.profiles import GraphMethodProfile, LIGHT

_METHODS: dict[str, "MethodDefinition"] = {}


@dataclass(frozen=True)
class MethodDefinition:
    key: str
    version: str
    label: str
    description: str
    default_lifecycle_steps: tuple[str, ...]
    next_action_strategy_key: str = "default"
    data_slices: frozenset[str] = frozenset()
    compatible_archetype_keys: frozenset[str] | Literal["*"] = "*"
    steering_elements: frozenset[str] = frozenset()
    ui_features: frozenset[str] = frozenset()
    graph_profile: GraphMethodProfile = LIGHT
    method_role: Literal["primary", "modifier"] = "primary"
    composes_with: frozenset[str] = frozenset()
    slot_map: tuple[tuple[str, str], ...] = ()


def ui_features_as_dict(features: frozenset[str]) -> dict[str, bool]:
    return {key: True for key in sorted(features)}


def method_compatible_with_archetype(
    method: MethodDefinition,
    archetype_key: str,
    *,
    om_capabilities: frozenset[str] | None = None,
) -> bool:
    compat = method.compatible_archetype_keys
    if compat != "*" and archetype_key not in compat:
        return False
    if om_capabilities is not None and method.data_slices:
        if not method.data_slices.issubset(om_capabilities):
            return False
    return True


def list_compatible_methods(
    archetype_key: str,
    *,
    om_capabilities: frozenset[str] | None = None,
) -> tuple[MethodDefinition, ...]:
    return tuple(
        m
        for m in _METHODS.values()
        if method_compatible_with_archetype(
            m, archetype_key, om_capabilities=om_capabilities
        )
    )


def register_method(defn: MethodDefinition) -> None:
    if defn.key in _METHODS:
        raise ValueError(f"Method already registered: {defn.key}")
    if defn.steering_elements:
        validate_steering_element_keys(defn.steering_elements)
    if defn.method_role == "modifier" and not defn.composes_with:
        raise ValueError(
            f"Modifier method {defn.key} must declare composes_with primary keys"
        )
    if defn.slot_map:
        from steering.lifecycle.slot_maps import slot_map_from_tuple

        slot_map_from_tuple(defn.slot_map)
    _METHODS[defn.key] = defn


def validate_method_registry() -> None:
    """Post-bootstrap: strategy keys exist; composes_with targets are primary methods."""
    from steering.strategies.next_action.registry import get_next_action_strategy

    for method in _METHODS.values():
        strategy_key = method.next_action_strategy_key
        if strategy_key != "default" and not get_next_action_strategy(strategy_key):
            raise ValueError(
                f"Method {method.key} references unknown next_action_strategy: {strategy_key}"
            )
        for primary_key in method.composes_with:
            primary = get_method(primary_key)
            if not primary:
                raise ValueError(
                    f"Method {method.key} composes_with unknown method: {primary_key}"
                )
            if primary.method_role != "primary":
                raise ValueError(
                    f"Method {method.key} composes_with non-primary method: {primary_key}"
                )


def list_composable_modifiers(primary_method_key: str) -> tuple[MethodDefinition, ...]:
    """Modifier methods that may compose with the given primary method."""
    return tuple(
        m
        for m in _METHODS.values()
        if m.method_role == "modifier" and primary_method_key in m.composes_with
    )


def get_method(key: str) -> MethodDefinition | None:
    return _METHODS.get(key)


def list_methods() -> tuple[MethodDefinition, ...]:
    return tuple(_METHODS.values())


def method_to_dict(
    method: MethodDefinition,
    *,
    include_compatibility: bool = True,
) -> dict[str, Any]:
    payload: dict[str, Any] = {
        "key": method.key,
        "version": method.version,
        "label": method.label,
        "description": method.description,
        "next_action_strategy_key": method.next_action_strategy_key,
        "steering_elements": sorted(method.steering_elements),
        "ui_features": ui_features_as_dict(method.ui_features),
        "graph_profile": {
            "enforce_gate_blocking": method.graph_profile.enforce_gate_blocking,
            "emphasize_fulfillment": method.graph_profile.emphasize_fulfillment,
        },
        "method_role": method.method_role,
        "composes_with": sorted(method.composes_with),
        "data_slices": sorted(method.data_slices),
    }
    if method.slot_map:
        payload["slot_map"] = {
            step: status for step, status in method.slot_map
        }
    if include_compatibility:
        compat = method.compatible_archetype_keys
        payload["compatible_archetype_keys"] = (
            "*" if compat == "*" else sorted(compat)
        )
    return payload


def clear_methods_for_tests() -> None:
    _METHODS.clear()