from __future__ import annotations

from typing import Iterable

from flownet.engine.operators.base import OperatorHandler
from flownet.engine.operators.defaults import register_default_operators


RESERVED_PREFIX = "sys/"


class OperatorRegistry:
    def __init__(self, default_handler: OperatorHandler | None = None):
        self._default_handler = default_handler or OperatorHandler()
        self._handlers: dict[str, OperatorHandler] = {}

    @property
    def default_handler(self) -> OperatorHandler:
        return self._default_handler

    def register(self, operator_key: str, handler: OperatorHandler) -> None:
        if not operator_key:
            raise ValueError("operator_key must be a non-empty string")
        if operator_key.startswith(RESERVED_PREFIX):
            raise ValueError(
                f"operator_key '{operator_key}' uses reserved prefix '{RESERVED_PREFIX}'."
            )
        self._handlers[operator_key] = handler

    def register_internal(self, operator_key: str, handler: OperatorHandler) -> None:
        if not operator_key:
            raise ValueError("operator_key must be a non-empty string")
        self._handlers[operator_key] = handler

    def resolve(self, operator_key: str | None) -> OperatorHandler:
        if not operator_key:
            return self._default_handler
        return self._handlers.get(operator_key, self._default_handler)

    def keys(self) -> Iterable[str]:
        return self._handlers.keys()

    def stats_snapshot(self) -> dict[str, dict[str, int]]:
        snapshot: dict[str, dict[str, int]] = {}
        for key, handler in self._handlers.items():
            stats_fn = getattr(handler, "stats_snapshot", None)
            if not callable(stats_fn):
                continue
            stats = stats_fn()
            if stats:
                snapshot[key] = stats
        return snapshot

    @classmethod
    def with_defaults(cls) -> "OperatorRegistry":
        registry = cls()
        return register_default_operators(registry)


__all__ = ["OperatorRegistry"]
