from __future__ import annotations
import asyncio
import contextlib
import time
import uuid
from dataclasses import dataclass
from typing import Awaitable, Callable, Dict, List, Optional, Set, Tuple

from .events import BusEvent
from .context import ContextObject
from .scheduler import PriorityScheduler
from .routing import ConfidenceRouter
from ..memory.store import EventStore, InMemoryEventStore

SubscriberFn = Callable[[BusEvent], Awaitable[None]]
FilterFn = Callable[[BusEvent], bool]

@dataclass(frozen=True)
class Subscription:
    agent_id: str
    filter_fn: FilterFn
    callback: SubscriberFn
    min_confidence: float = 0.0
    topics: Optional[Set[str]] = None

class CognitiveBus:
    """Async blackboard bus with:
    - priority scheduling + backpressure
    - policy routing (confidence-based)
    - optional persistence store for replay/checkpoints
    """

    def __init__(
        self,
        max_queue: int = 10_000,
        router: Optional[ConfidenceRouter] = None,
        store: Optional[EventStore] = None,
    ):
        self._subs: List[Subscription] = []
        self._running = False

        self._scheduler = PriorityScheduler(maxsize=max_queue)
        self._scheduler.set_handler(self._dispatch)

        self._store_mem: Dict[str, BusEvent] = {}
        self.router = router or ConfidenceRouter()
        self.store: EventStore = store or InMemoryEventStore()

    def snapshot(self) -> List[BusEvent]:
        return list(self._store_mem.values())

    def get(self, event_id: str) -> Optional[BusEvent]:
        return self._store_mem.get(event_id)

    def subscribe(
        self,
        *,
        agent_id: str,
        callback: SubscriberFn,
        filter_fn: Optional[FilterFn] = None,
        topics: Optional[Set[str]] = None,
        min_confidence: float = 0.0,
    ) -> None:
        if filter_fn is None:
            filter_fn = lambda e: True
        self._subs.append(Subscription(agent_id, filter_fn, callback, min_confidence, topics))

    async def publish(self, event: BusEvent) -> None:
        # apply routing policy before persistence/dispatch
        decision = self.router.decide(event)
        if decision.priority is not None:
            event.priority = decision.priority
        if decision.tags_add:
            event.tags = list(dict.fromkeys([*(event.tags or []), *decision.tags_add]))
        if decision.intended_for is not None:
            event.intended_for = decision.intended_for

        self._store_mem[event.event_id] = event
        await self.store.put_event(event)
        await self._scheduler.put(event)

    async def emit(
        self,
        *,
        topic: str,
        type: str,
        sender: str,
        context: Optional[ContextObject] = None,
        parent_id: Optional[str] = None,
        intended_for: Optional[List[str]] = None,
        tags: Optional[List[str]] = None,
        confidence: Optional[float] = None,
        priority: int = 5,
        meta: Optional[Dict] = None,
    ) -> BusEvent:
        event = BusEvent(
            event_id=str(uuid.uuid4()),
            topic=topic,
            type=type,
            sender=sender,
            parent_id=parent_id,
            intended_for=intended_for or [],
            tags=tags or [],
            confidence=confidence,
            priority=priority,
            meta=meta or {"ts": time.time()},
            context=context or ContextObject(),
        )
        await self.publish(event)
        return event

    async def start(self) -> None:
        if self._running:
            return
        self._running = True
        await self._scheduler.start()

    async def stop(self) -> None:
        self._running = False
        await self._scheduler.stop()

    async def replay_recent(self, limit: int = 200) -> List[BusEvent]:
        """Replay last N events from the configured EventStore into in-memory blackboard (no dispatch)."""
        events = await self.store.list_recent(limit=limit)
        for ev in events:
            self._store_mem[ev.event_id] = ev
        return events

    async def _dispatch(self, event: BusEvent) -> None:
        direct = set(event.intended_for or [])
        targets: List[Tuple[float, Subscription]] = []

        for sub in self._subs:
            if sub.topics and event.topic not in sub.topics:
                continue
            if direct and sub.agent_id not in direct:
                continue
            if event.confidence is not None and event.confidence < sub.min_confidence:
                continue
            if sub.filter_fn(event):
                rank = float(event.confidence or 0.0)
                targets.append((rank, sub))

        targets.sort(key=lambda x: x[0], reverse=True)
        await asyncio.gather(*(sub.callback(event) for _, sub in targets), return_exceptions=True)
