"""Attention contributor registry — K-Ext-2."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Callable

from steering.attention.context import AttentionEvalContext

AttentionContribute = Callable[[AttentionEvalContext], list[dict]]


@dataclass(frozen=True)
class AttentionContributor:
    key: str
    requires_elements: frozenset[str] = field(default_factory=frozenset)
    requires_any_element: frozenset[str] = field(default_factory=frozenset)
    requires_primary_methods: frozenset[str] = field(default_factory=frozenset)
    contribute: AttentionContribute = lambda _ctx: []


_CONTRIBUTORS: list[AttentionContributor] = []


def register_attention_contributor(contributor: AttentionContributor) -> None:
    _CONTRIBUTORS.append(contributor)


def collect_attention_items(ctx: AttentionEvalContext) -> list[dict]:
    elements = ctx.steering_elements
    primary = ctx.binding.primary_method_key
    items: list[dict] = []

    for contributor in _CONTRIBUTORS:
        if (
            contributor.requires_elements
            and not contributor.requires_elements <= elements
        ):
            continue
        if contributor.requires_any_element and not (
            contributor.requires_any_element & elements
        ):
            continue
        if (
            contributor.requires_primary_methods
            and primary not in contributor.requires_primary_methods
        ):
            continue
        items.extend(contributor.contribute(ctx))

    for item in items:
        item.setdefault("data_source", "steering_kernel")

    items.sort(
        key=lambda x: {"critical": 0, "warning": 1, "info": 2}.get(
            x.get("severity", "info"), 9
        )
    )
    return items