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

DESIGN §7 简化(0004 一并)—— 单活 run 形态下 run_id 是冗余的(同 task 同时
最多 1 个活 run);broker 内部全用 task_id 索引,客户端只需要 task_id 即可
订阅 / cancel,不再需要先拿 run_id。

两个实现,env `ZCBOT_REDIS_URL` 决定(import 时定,进程生命周期不变):
- **LocalRunBroker**(默认,env 不设):进程内 asyncio queue,零依赖 —— 本地 dev /
  单实例部署用。多进程不共享:蓝绿并存期跨实例订阅/cancel 送达不到(RUN.md B 档边缘)。
- **RedisRunBroker**(env 设了):event 走 Redis PUB/SUB(`zcbot:ev:<task_id>`),
  done/cancel 走带 TTL 的 key —— 蓝绿双实例跨进程可达,也是稳态双实例分流的前提
  (DESIGN §7.0)。启动 ping 不通 fail-fast(同 sandbox init 纪律);运行中 publish
  失败 log + 丢帧(直播降级,run 不受影响)。

公共语义(两实现一致):
- emit() 从工作线程调(agent.run 在 to_thread 跑);SSE generator await queue.get()。
- 同一 task 多个订阅者(刷新页面 / 多 tab)— 每个订阅 1 个独立 queue。
- run 结束 → close(task_id) 给所有订阅者派 done;晚到的订阅者立即收 done 断流
  (local 用 _done set;redis 用 `zcbot:done:<tid>` 60s TTL key —— 主保护是 SSE
  端点先查 DB run_status,key 只兜"查完 DB 到订阅完成之间 run 恰好收尾"的窗口)。
- 同 task 起新 run → start(task_id) 清 done 标记。
- cancel:request_cancel 置位,AgentLoop 在 chunk 间 poll is_cancelled() 即退;
  redis 版是 `zcbot:cancel:<tid>` key(GET ~0.1ms/次,localhost 可忽略),跨实例可达。
- 不持久化 event — messages 已落 PG,刷新走 GET /messages 看历史;订阅建立前
  emit 的帧两个实现都会丢(设计选择)。
"""
from __future__ import annotations

import asyncio
import json
import os
import threading
import time
from collections import defaultdict
from typing import Any, Optional
from uuid import UUID


class LocalRunBroker:
    """进程内实现(原 RunBroker)。

    线程模型:
    - bind_loop(loop) 在 FastAPI startup 调一次,记录 asyncio loop 引用。
    - emit() 调用方可能在任意线程;put_nowait 是 thread-unsafe(asyncio.Queue 设计
      前提是单 loop),所以走 call_soon_threadsafe 跨回 loop 线程再 put。
    """

    def __init__(self) -> None:
        self._subs: dict[UUID, set[asyncio.Queue]] = defaultdict(set)
        # 已经发完 done 的 task — 后来订阅者直接收到 done,避免无限等
        self._done: set[UUID] = set()
        # cancel signal per-task。AgentLoop 在 BG 线程里 poll is_cancelled() 决定是否退;
        # request_cancel 可在 BG 还没 register 时调用(setdefault),BG 启动后第一次
        # check 即看到。run 完成在 finally 里 clear_cancel 回收。
        self._cancel_flags: dict[UUID, threading.Event] = {}
        self._loop: Optional[asyncio.AbstractEventLoop] = None

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

    async def shutdown(self) -> None:
        """lifespan finally 调;local 无外部资源,no-op。"""

    def start(self, task_id: UUID) -> None:
        """同 task 起新 run 时调:清 _done 标记,让新订阅者能看到流式。
        cancel flag 在 finally 里 clear_cancel 清,这里不动(避免擦掉刚刚 request_cancel 的请求)。
        """
        self._done.discard(task_id)

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

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

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

    def emit(self, task_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(task_id, [])):
            loop.call_soon_threadsafe(q.put_nowait, event)

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

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

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

    def total_subscribers(self) -> int:
        """本进程 SSE 订阅者总数(所有 task 累加),供 §8.4 并发监控用。"""
        return sum(len(s) for s in self._subs.values())

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

    # ─────────────── cancel signaling ───────────────

    def request_cancel(self, task_id: UUID) -> None:
        """主线程(HTTP handler)发的 cancel 信号 — BG 线程 poll is_cancelled() 看见即退。
        setdefault:即便 BG 还没注册 flag 也能 set,BG 启动后第一次 check 立刻看见。
        """
        self._cancel_flags.setdefault(task_id, threading.Event()).set()

    def is_cancelled(self, task_id: UUID) -> bool:
        ev = self._cancel_flags.get(task_id)
        return bool(ev and ev.is_set())

    def clear_cancel(self, task_id: UUID) -> None:
        """run 真正退出(BG finally)清掉 flag,避免 dict 无限增长。"""
        self._cancel_flags.pop(task_id, None)


class RedisRunBroker:
    """Redis pub/sub 实现(蓝绿双实例跨进程,DESIGN §7.0)。

    - event:`PUBLISH zcbot:ev:<task_id>`(JSON);单条 async pubsub reader 任务
      把消息路由到本进程各订阅 queue(订阅表 _subs 仍在本地 —— fan-out 的"最后
      一公里"还是进程内,Redis 只负责跨进程那一跳)。
    - done:close 先 `SETEX zcbot:done:<tid> 60`再 publish done;subscribe 先挂
      channel 再查 key,晚到订阅者不漏。
    - cancel:`SETEX zcbot:cancel:<tid> 3600` / GET / DEL,跨实例可达。
    - 容错:bind_loop ping 不通直接 raise(fail-fast);运行中 redis 操作失败
      log(10s 节流)+ 降级(emit 丢帧、is_cancelled 返 False),run 不受影响;
      reader 断线 1s 退避重连并重挂全部活跃 channel。
    """

    _EV_PREFIX = "zcbot:ev:"
    _DONE_PREFIX = "zcbot:done:"
    _CANCEL_PREFIX = "zcbot:cancel:"
    _CTL_CHANNEL = "zcbot:ev:__ctl__"   # 常挂空 channel,保 pubsub 始终有订阅可 get_message
    _DONE_TTL = 60
    _CANCEL_TTL = 3600

    def __init__(self, url: str, *, sync_client=None, async_client=None) -> None:
        import redis as _redis
        import redis.asyncio as _aredis
        self._url = url
        # sync client:emit/close/cancel 从工作线程调(redis-py 自带连接池,线程安全);
        # start/request_cancel 也会在 loop 线程上调 —— localhost 单次 ~0.1ms,超时收紧
        # 到 2s 防 redis 挂掉时把 event loop 卡死太久。
        self._redis = sync_client if sync_client is not None else _redis.Redis.from_url(
            url, decode_responses=True, socket_timeout=2.0, socket_connect_timeout=2.0,
        )
        self._aredis = async_client if async_client is not None else _aredis.from_url(
            url, decode_responses=True,
        )
        self._pubsub = None
        self._reader_task: Optional[asyncio.Task] = None
        self._subs: dict[UUID, set[asyncio.Queue]] = defaultdict(set)
        self._loop: Optional[asyncio.AbstractEventLoop] = None
        self._last_err_log = 0.0

    # ─────────────── helpers ───────────────

    def _chan(self, task_id: UUID) -> str:
        return f"{self._EV_PREFIX}{task_id}"

    def _log_err(self, what: str, e: Exception) -> None:
        """redis 故障时 emit/is_cancelled 每 chunk 都会撞,10s 节流防刷屏。"""
        now = time.monotonic()
        if now - self._last_err_log >= 10:
            self._last_err_log = now
            print(f"[broker] redis {what} failed: {type(e).__name__}: {e} (降级继续,10s 内同类不再报)")

    # ─────────────── lifecycle ───────────────

    def bind_loop(self, loop: asyncio.AbstractEventLoop) -> None:
        """FastAPI startup 调一次:ping 验活(不通 fail-fast)+ 起 pubsub reader。"""
        self._loop = loop
        try:
            self._redis.ping()
        except Exception as e:
            raise RuntimeError(
                f"ZCBOT_REDIS_URL 已设但 Redis 连不上({type(e).__name__}: {e})——"
                f"确认 redis-server 在跑 / URL 对,或删掉该 env 回退进程内 broker"
            )
        self._reader_task = loop.create_task(self._reader(), name="redis-broker-reader")
        host = self._url.rsplit("@", 1)[-1]  # 不把密码打进日志
        print(f"[broker] redis pub/sub enabled ({host})")

    async def shutdown(self) -> None:
        """lifespan finally 调:停 reader、断两个客户端。"""
        if self._reader_task is not None:
            self._reader_task.cancel()
            try:
                await self._reader_task
            except (asyncio.CancelledError, Exception):
                pass
        if self._pubsub is not None:
            try:
                await self._pubsub.aclose()
            except Exception:
                pass
        try:
            await self._aredis.aclose()
        except Exception:
            pass
        try:
            self._redis.close()
        except Exception:
            pass

    async def _reader(self) -> None:
        """常驻:单条 pubsub 连接收所有活跃 channel 的消息,路由到本地订阅 queue。
        断线 1s 退避重连,并把 _subs 里所有活跃 channel 重挂上(订阅状态不丢)。"""
        while True:
            try:
                self._pubsub = self._aredis.pubsub()
                channels = [self._CTL_CHANNEL] + [self._chan(t) for t in self._subs]
                await self._pubsub.subscribe(*channels)
                while True:
                    msg = await self._pubsub.get_message(
                        ignore_subscribe_messages=True, timeout=5.0
                    )
                    if msg is None or msg.get("type") != "message":
                        continue
                    ch = msg.get("channel") or ""
                    if ch == self._CTL_CHANNEL:
                        continue
                    try:
                        tid = UUID(ch[len(self._EV_PREFIX):])
                        event = json.loads(msg["data"])
                    except (ValueError, TypeError):
                        continue  # 不认识的 channel / 坏 payload,丢
                    for q in list(self._subs.get(tid, ())):
                        q.put_nowait(event)
            except asyncio.CancelledError:
                raise
            except Exception as e:
                self._log_err("reader", e)
                try:
                    if self._pubsub is not None:
                        await self._pubsub.aclose()
                except Exception:
                    pass
                await asyncio.sleep(1)

    # ─────────────── pub/sub ───────────────

    def start(self, task_id: UUID) -> None:
        try:
            self._redis.delete(f"{self._DONE_PREFIX}{task_id}")
        except Exception as e:
            self._log_err("start/del-done", e)

    def subscribe(self, task_id: UUID) -> asyncio.Queue:
        """调用方:SSE handler(loop 线程)。queue 立即返回;channel 挂载与 done-key
        补查异步完成 —— 挂载完成前 publish 的帧会丢(与 local 版"订阅前的帧丢"同类,
        SSE 端点已先查过 DB run_status 兜大头)。"""
        q: asyncio.Queue = asyncio.Queue()
        first = not self._subs[task_id]
        self._subs[task_id].add(q)

        async def _setup() -> None:
            try:
                if first and self._pubsub is not None:
                    await self._pubsub.subscribe(self._chan(task_id))
                # 先挂 channel 再查 done key:close() 是先 SETEX 再 publish,
                # 两边一交叉,晚到订阅者要么收到 publish、要么查到 key,不漏
                if await self._aredis.exists(f"{self._DONE_PREFIX}{task_id}"):
                    q.put_nowait({"type": "done"})
            except Exception as e:
                self._log_err("subscribe-setup", e)
                q.put_nowait({"type": "done"})  # 保守断流,客户端自动重连

        if self._loop is not None:
            self._loop.create_task(_setup())
        return q

    def unsubscribe(self, task_id: UUID, q: asyncio.Queue) -> None:
        subs = self._subs.get(task_id)
        if subs is None:
            return
        subs.discard(q)
        if not subs:
            del self._subs[task_id]
            ch = self._chan(task_id)

            async def _unsub() -> None:
                try:
                    if self._pubsub is not None:
                        await self._pubsub.unsubscribe(ch)
                except Exception:
                    pass  # reader 重连时按 _subs 重挂,残留订阅无害

            if self._loop is not None:
                self._loop.create_task(_unsub())

    def emit(self, task_id: UUID, event: dict[str, Any]) -> None:
        try:
            self._redis.publish(
                self._chan(task_id), json.dumps(event, ensure_ascii=False, default=str)
            )
        except Exception as e:
            self._log_err("publish", e)  # 丢帧降级,run 不受影响

    def close(self, task_id: UUID) -> None:
        try:
            self._redis.set(f"{self._DONE_PREFIX}{task_id}", "1", ex=self._DONE_TTL)
        except Exception as e:
            self._log_err("close/set-done", e)
        self.emit(task_id, {"type": "done"})

    def n_subscribers(self, task_id: UUID) -> int:
        return len(self._subs.get(task_id, set()))

    def total_subscribers(self) -> int:
        """本进程的订阅数(蓝绿双实例各报各的,admin 面板看单实例)。"""
        return sum(len(s) for s in self._subs.values())

    def is_done(self, task_id: UUID) -> bool:
        try:
            return bool(self._redis.exists(f"{self._DONE_PREFIX}{task_id}"))
        except Exception as e:
            self._log_err("is_done", e)
            return False

    # ─────────────── cancel signaling ───────────────

    def request_cancel(self, task_id: UUID) -> None:
        try:
            self._redis.set(f"{self._CANCEL_PREFIX}{task_id}", "1", ex=self._CANCEL_TTL)
        except Exception as e:
            self._log_err("request_cancel", e)

    def is_cancelled(self, task_id: UUID) -> bool:
        """AgentLoop 在 stream chunk 间 poll(工作线程);redis 挂了返 False 降级
        (cancel 暂不可达,等价于 redis 前的跨实例行为,run 不受影响)。"""
        try:
            return self._redis.get(f"{self._CANCEL_PREFIX}{task_id}") is not None
        except Exception as e:
            self._log_err("is_cancelled", e)
            return False

    def clear_cancel(self, task_id: UUID) -> None:
        try:
            self._redis.delete(f"{self._CANCEL_PREFIX}{task_id}")
        except Exception as e:
            self._log_err("clear_cancel", e)


def _make_broker():
    """env `ZCBOT_REDIS_URL` 在 → Redis 版(蓝绿跨进程);不在 → 进程内版(dev/单实例)。
    import 时定一次,与 ZCBOT_INSTANCE 同款读法。"""
    url = os.getenv("ZCBOT_REDIS_URL", "").strip()
    if url:
        return RedisRunBroker(url)
    return LocalRunBroker()


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