space/manager.py

from __future__ import annotations

import asyncio
from collections import deque
from collections.abc import Iterable

from pydantic import BaseModel

from .event import EventEnvelope, InterruptData, PluginErrorData
from .plugin import EventHandler, Plugin


class PluginManager:
    def __init__(self, *, event_limit: int = 1_000) -> None:
        if event_limit < 0:
            raise ValueError("event limit must not be negative")
        self._plugins: dict[str, Plugin] = {}
        self._subscribers: dict[str, list[tuple[Plugin, EventHandler]]] = {}
        self._contracts: dict[str, type[BaseModel]] = {
            "error": PluginErrorData,
            "interrupt": InterruptData,
        }
        self._events: deque[EventEnvelope[BaseModel]] = deque(maxlen=event_limit)
        self._dispatch_tasks: set[asyncio.Task[None]] = set()

    def register_event(self, name: str, contract: type[BaseModel]) -> None:
        if not name or not issubclass(contract, BaseModel):
            raise ValueError("event registrations require a name and Pydantic contract")
        self._contracts[name] = contract

    async def enable_plugin(self, plugin: Plugin) -> None:
        if plugin.name in self._plugins and self._plugins[plugin.name] is not plugin:
            raise ValueError(f"plugin already registered: {plugin.name}")
        if plugin.enabled:
            return
        plugin.manager = self
        if not plugin._initialized:
            await plugin.initialize()
            plugin._initialized = True
        plugin.enabled = True
        self._plugins[plugin.name] = plugin
        for name, handler in plugin.iter_event_handlers().items():
            self._subscribers.setdefault(name, []).append((plugin, handler))
        plugin._start_worker()

    async def disable_plugin(self, name: str) -> None:
        plugin = self._plugins.get(name)
        if plugin is None or not plugin.enabled:
            return
        plugin.enabled = False
        self._subscribers = {
            event: [(item, handler) for item, handler in subscribers if item is not plugin]
            for event, subscribers in self._subscribers.items()
        }
        await plugin._stop_worker()

    async def close(self) -> None:
        for name in tuple(self._plugins):
            await self.disable_plugin(name)

    def get_plugins(self) -> tuple[Plugin, ...]:
        return tuple(self._plugins.values())

    def get_events(self) -> tuple[EventEnvelope[BaseModel], ...]:
        return tuple(self._events)

    async def emit(
        self,
        name: str,
        data: BaseModel,
        *,
        source: str = "space",
        target: str | None = None,
        capabilities: Iterable[str] = (),
        only_once: bool = False,
    ) -> EventEnvelope[BaseModel]:
        contract = self._contracts.get(name)
        if contract is not None and not isinstance(data, contract):
            raise TypeError(f"{name} requires {contract.__name__}")
        required = tuple(capabilities)
        envelope = EventEnvelope[BaseModel](
            name=name,
            data=data,
            source=source,
            target=target,
            capabilities=list(required) or None,
            only_once=only_once,
        )
        self._events.append(envelope)
        recipients = [
            (plugin, handler)
            for plugin, handler in self._subscribers.get(name, [])
            if plugin.enabled
            and (plugin.receive_self_events or plugin.name != source.split("/", maxsplit=1)[0])
            and (target is None or plugin.name == target)
        ]
        if only_once:
            recipients = recipients[:1]
        for plugin, handler in recipients:
            task = asyncio.create_task(plugin._enqueue(envelope, handler))
            self._dispatch_tasks.add(task)
            task.add_done_callback(self._dispatch_tasks.discard)
        return envelope

    async def wait_for_idle(self, *, timeout: float = 5) -> None:
        async with asyncio.timeout(timeout):
            while self._dispatch_tasks or any(not plugin.queue.empty() for plugin in self._plugins.values()):
                if self._dispatch_tasks:
                    await asyncio.gather(*tuple(self._dispatch_tasks))
                await asyncio.gather(*(plugin.queue.join() for plugin in self._plugins.values() if plugin.enabled))
                await asyncio.sleep(0)

    async def report_plugin_error(self, plugin: Plugin, envelope: EventEnvelope[BaseModel], error: Exception) -> None:
        if envelope.name != "error":
            await self.emit(
                "error",
                PluginErrorData(
                    plugin=plugin.name,
                    event_name=envelope.name,
                    error_type=type(error).__name__,
                    message=str(error),
                ),
                source=plugin.name,
            )