PMFlow/backend/app/core/events.py

"""Neural Sync - Event Bus System。"""

from __future__ import annotations

import asyncio
import logging
from contextlib import asynccontextmanager
from typing import Any, Callable, Optional

from .events_models import Event, EventType

logger = logging.getLogger(__name__)


class EventBus:
    """事件总线核心。"""

    def __init__(self):
        self._subscribers: dict[EventType, list[Callable]] = {}
        self._wildcard_subscribers: list[Callable] = []
        self._event_history: list[Event] = []
        self._max_history = 1000
        self._metrics = {
            "published": 0,
            "delivered": 0,
            "failed": 0,
        }
        self._lock = asyncio.Lock()
        self._running = False

    @property
    def metrics(self) -> dict[str, int]:
        return self._metrics.copy()

    @property
    def history(self) -> list[dict[str, Any]]:
        return [event.to_dict() for event in self._event_history]

    async def start(self):
        self._running = True
        logger.info("EventBus started")

    async def stop(self):
        self._running = False
        logger.info("EventBus stopped")

    def subscribe(
        self,
        event_type: EventType,
        handler: Callable,
        filter_func: Optional[Callable[[Event], bool]] = None,
    ) -> Callable:
        async def wrapped_handler(event: Event):
            if filter_func and not filter_func(event):
                return
            try:
                await handler(event)
                self._metrics["delivered"] += 1
            except Exception as exc:
                logger.error(f"Error in event handler for {event_type.value}: {exc}", exc_info=True)
                self._metrics["failed"] += 1

        wrapped_handler.original = handler  # type: ignore[attr-defined]

        async def _subscribe():
            async with self._lock:
                if event_type not in self._subscribers:
                    self._subscribers[event_type] = []
                self._subscribers[event_type].append(wrapped_handler)

        asyncio.create_task(_subscribe())

        def unsubscribe():
            asyncio.create_task(self._unsubscribe(event_type, wrapped_handler))

        return unsubscribe

    def subscribe_wildcard(
        self,
        handler: Callable,
        filter_func: Optional[Callable[[Event], bool]] = None,
    ) -> Callable:
        async def wrapped_handler(event: Event):
            if filter_func and not filter_func(event):
                return
            try:
                await handler(event)
                self._metrics["delivered"] += 1
            except Exception as exc:
                logger.error(f"Error in wildcard event handler: {exc}", exc_info=True)
                self._metrics["failed"] += 1

        wrapped_handler.original = handler  # type: ignore[attr-defined]

        async def _subscribe():
            async with self._lock:
                self._wildcard_subscribers.append(wrapped_handler)

        asyncio.create_task(_subscribe())

        def unsubscribe():
            asyncio.create_task(self._unsubscribe_wildcard(wrapped_handler))

        return unsubscribe

    async def _unsubscribe(self, event_type: EventType, handler: Callable):
        async with self._lock:
            if event_type in self._subscribers:
                self._subscribers[event_type] = [
                    subscribed
                    for subscribed in self._subscribers[event_type]
                    if subscribed != handler
                ]

    async def _unsubscribe_wildcard(self, handler: Callable):
        async with self._lock:
            self._wildcard_subscribers = [
                subscribed
                for subscribed in self._wildcard_subscribers
                if subscribed != handler
            ]

    async def publish(self, event: Event) -> bool:
        if not self._running:
            logger.warning(f"EventBus not running, discarding event: {event.type.value}")
            return False

        self._metrics["published"] += 1

        self._event_history.append(event)
        if len(self._event_history) > self._max_history:
            self._event_history.pop(0)

        tasks = []
        if event.type in self._subscribers:
            for handler in self._subscribers[event.type]:
                tasks.append(handler(event))

        for handler in self._wildcard_subscribers:
            tasks.append(handler(event))

        if tasks:
            results = await asyncio.gather(*tasks, return_exceptions=True)
            failed = sum(1 for result in results if isinstance(result, Exception))
            if failed > 0:
                self._metrics["failed"] += failed

        logger.debug(f"Published event {event.type.value} to {len(tasks)} subscribers")
        return True

    async def publish_and_wait(self, event: Event, timeout: float = 5.0) -> list[Any]:
        if not self._running:
            raise RuntimeError("EventBus not running")

        self._metrics["published"] += 1
        self._event_history.append(event)

        tasks = []
        if event.type in self._subscribers:
            tasks.extend(self._subscribers[event.type])

        for handler in self._wildcard_subscribers:
            tasks.append(handler(event))

        if not tasks:
            return []

        try:
            return await asyncio.wait_for(
                asyncio.gather(*tasks, return_exceptions=True),
                timeout=timeout,
            )
        except asyncio.TimeoutError:
            logger.warning(f"Event {event.type.value} processing timed out after {timeout}s")
            raise

    def clear_history(self):
        self._event_history.clear()

    def get_subscriber_count(self, event_type: Optional[EventType] = None) -> int:
        if event_type:
            return len(self._subscribers.get(event_type, []))
        return sum(len(subs) for subs in self._subscribers.values()) + len(self._wildcard_subscribers)


_global_event_bus: Optional[EventBus] = None


def get_event_bus() -> EventBus:
    global _global_event_bus
    if _global_event_bus is None:
        _global_event_bus = EventBus()
    return _global_event_bus


@asynccontextmanager
async def event_bus_context():
    bus = get_event_bus()
    await bus.start()
    try:
        yield bus
    finally:
        await bus.stop()


async def publish_event(event_type: EventType, data: dict[str, Any], **kwargs) -> bool:
    event = Event(type=event_type, data=data, **kwargs)
    return await get_event_bus().publish(event)


def subscribe_event(event_type: EventType, handler: Callable) -> Callable:
    return get_event_bus().subscribe(event_type, handler)


def subscribe_all_events(handler: Callable) -> Callable:
    return get_event_bus().subscribe_wildcard(handler)


__all__ = [
    "Event",
    "EventBus",
    "EventType",
    "event_bus_context",
    "get_event_bus",
    "publish_event",
    "subscribe_all_events",
    "subscribe_event",
]