from __future__ import annotations

import asyncio
from collections import defaultdict
from dataclasses import dataclass

from app.core.contracts import RuntimeEvent
from app.events.event_bus import EventBus


class StreamingManager:
    """Simple in-process projection from event bus to websocket consumers."""

    def __init__(self, event_bus: EventBus) -> None:
        self._event_bus = event_bus
        self._subscribers: dict[str, list[StreamSubscriber]] = defaultdict(list)
        self._event_bus.subscribe(self._on_event)

    def replay_events(self, task_id: str) -> list[RuntimeEvent]:
        return self._event_bus.list_for_task(task_id)

    def subscribe(self, task_id: str) -> asyncio.Queue[RuntimeEvent]:
        queue: asyncio.Queue[RuntimeEvent] = asyncio.Queue()
        self._subscribers[task_id].append(
            StreamSubscriber(loop=asyncio.get_running_loop(), queue=queue)
        )
        return queue

    def unsubscribe(self, task_id: str, queue: asyncio.Queue[RuntimeEvent]) -> None:
        listeners = self._subscribers.get(task_id, [])
        for listener in list(listeners):
            if listener.queue is queue:
                listeners.remove(listener)
                break
        if not listeners and task_id in self._subscribers:
            del self._subscribers[task_id]

    def _on_event(self, event: RuntimeEvent) -> None:
        for listener in list(self._subscribers.get(event.task_id, [])):
            listener.loop.call_soon_threadsafe(listener.queue.put_nowait, event)


@dataclass
class StreamSubscriber:
    loop: asyncio.AbstractEventLoop
    queue: asyncio.Queue[RuntimeEvent]