"""
SSEManager — pub/sub hub for Server-Sent Events.

Task runner calls publish() → all connected SSE clients watching that task_id receive the event.

Replay buffer: each task_id stores the last REPLAY_MAXLEN events so that clients who connect
slightly after publishing starts (race between POST /discover and GET /events) still receive
the initial events (disc_server, disc_rack_start, etc.).
"""

import asyncio
import json
from collections import deque
from dataclasses import dataclass, field
from datetime import datetime, timezone

REPLAY_MAXLEN = 200


@dataclass
class SSEEvent:
    type: str        # "task_event" | "task_status" | "task_result"
    task_id: str
    data: dict

    def encode(self) -> str:
        payload = json.dumps({"task_id": self.task_id, **self.data})
        return f"event: {self.type}\ndata: {payload}\n\n"


class SSEManager:
    def __init__(self) -> None:
        # task_id -> list[Queue]
        self._subscribers: dict[str, list[asyncio.Queue]] = {}
        # task_id -> deque of recent events for late-joining subscribers
        self._replay: dict[str, deque[SSEEvent]] = {}
        self._lock = asyncio.Lock()

    def clear_replay(self, task_id: str) -> None:
        """Call before starting a new task to discard stale events from a previous run."""
        self._replay.pop(task_id, None)

    async def subscribe(self, task_id: str, queue: asyncio.Queue) -> None:
        async with self._lock:
            self._subscribers.setdefault(task_id, []).append(queue)
            # Replay buffered events so late subscribers catch up
            for event in self._replay.get(task_id, []):
                try:
                    queue.put_nowait(event)
                except asyncio.QueueFull:
                    break

    async def unsubscribe(self, task_id: str, queue: asyncio.Queue) -> None:
        async with self._lock:
            subs = self._subscribers.get(task_id, [])
            try:
                subs.remove(queue)
            except ValueError:
                pass
            if not subs:
                self._subscribers.pop(task_id, None)

    async def publish(self, event: SSEEvent) -> None:
        async with self._lock:
            buf = self._replay.setdefault(event.task_id, deque(maxlen=REPLAY_MAXLEN))
            buf.append(event)
            queues = list(self._subscribers.get(event.task_id, []))
        for queue in queues:
            try:
                queue.put_nowait(event)
            except asyncio.QueueFull:
                pass  # slow client, drop event

    def subscriber_count(self, task_id: str) -> int:
        return len(self._subscribers.get(task_id, []))


# Module-level singleton shared across all requests
sse_manager = SSEManager()