"""RunBroker:in-process pub/sub,把 agent run 产生的 event fan-out 给所有 SSE 订阅者。

设计:
- emit() 从工作线程调(agent.run 在 to_thread 跑),用 loop.call_soon_threadsafe
  桥到 asyncio queue;SSE generator await queue.get() 拉出来推流。
- 同一 run_id 多个订阅者(刷新页面 / 多 tab / 桌面+移动)— 每个订阅 1 个独立 queue。
- run 结束 → broker.close(run_id) 给所有订阅者派一条 done;新订阅者(在 done 后到的)
  立即收到 done 并断流(不漏不挂)。
- 进程内单实例 / 多进程不共享 — 个人 SaaS 单 worker 够用;真要扩多 worker 再上 Redis。
- 不持久化 event — messages 已落 PG,刷新页面走 G3 静态视图能看历史;真要"刷新继续看
  实时流"未来加 event log 表 + backfill。

线程模型:
- broker.bind_loop(loop) 在 FastAPI startup 调一次,记录 asyncio loop 引用。
- emit() 调用方可能在任意线程;put_nowait 是 thread-unsafe(asyncio.Queue 设计前提
  是单 loop),所以走 call_soon_threadsafe 跨回 loop 线程再 put。
- subscribe / unsubscribe / close 也都用 call_soon_threadsafe 包,避免 race
  (实测 SSE generator 在 finally 里 unsubscribe,这个就在 loop 线程,直接调也行)。
"""
from __future__ import annotations

import asyncio
from collections import defaultdict
from typing import Any, Optional
from uuid import UUID


class RunBroker:
    def __init__(self) -> None:
        self._subs: dict[UUID, set[asyncio.Queue]] = defaultdict(set)
        # 已经发完 done 的 run — 后来订阅者直接收到 done,避免无限等
        self._done: set[UUID] = set()
        self._loop: Optional[asyncio.AbstractEventLoop] = None

    def bind_loop(self, loop: asyncio.AbstractEventLoop) -> None:
        """FastAPI startup 调一次。"""
        self._loop = loop

    def subscribe(self, run_id: UUID) -> asyncio.Queue:
        """订阅 run 的 event 流。已 done 的 run 立刻在 queue 放一条 done。

        调用方:SSE handler(在 asyncio loop 线程内)。
        """
        q: asyncio.Queue = asyncio.Queue()
        if run_id in self._done:
            q.put_nowait({"type": "done"})
        else:
            self._subs[run_id].add(q)
        return q

    def unsubscribe(self, run_id: UUID, q: asyncio.Queue) -> None:
        """SSE generator finally 清理。"""
        self._subs.get(run_id, set()).discard(q)
        if run_id in self._subs and not self._subs[run_id]:
            del self._subs[run_id]

    def emit(self, run_id: UUID, event: dict[str, Any]) -> None:
        """从工作线程调:把 event 推给所有订阅者。

        如果没人订阅(run 在跑但没浏览器连上),event 丢弃 — 这是设计选择
        (event 不持久化,messages 走 PG)。
        """
        loop = self._loop
        if loop is None:
            return  # 还没 bind,丢弃(测试 / 启动竞态)
        for q in list(self._subs.get(run_id, [])):
            loop.call_soon_threadsafe(q.put_nowait, event)

    def close(self, run_id: UUID) -> None:
        """run 结束:派 done 给所有订阅者,标记 run_id 为已完成。

        从工作线程调(agent.run 完成 / 抛异常 finally 清理)。
        """
        self.emit(run_id, {"type": "done"})
        self._done.add(run_id)
        # subs 不在这里立即删 — SSE generator 会先收到 done、yield 它、走到
        # finally unsubscribe;此处 emit 后立即删会让那次 emit 之后的清理无的放矢。

    def n_subscribers(self, run_id: UUID) -> int:
        """供测试 / 监控用。"""
        return len(self._subs.get(run_id, set()))

    def is_done(self, run_id: UUID) -> bool:
        return run_id in self._done


# 进程内单例 — FastAPI lifespan 里 bind_loop;agent / sink / SSE handler 共享。
broker = RunBroker()