"""SimScheduler:事件驱动的 agent 唤醒调度(去回合制的核心)。

每个 agent 只在三种时刻被唤醒(docs/agent-runtime-v2.md §1.1):
1. 预约时刻到(动作完成 / 思考节拍);2. 被中断(高显著度事件);3. 世界关停。
asyncio 单线程驱动:堆顶睡到点 → 并发触发 handler;中断插队即时生效。
"""

from __future__ import annotations

import asyncio
import heapq
import itertools

from genesis.obs.logging_setup import get_logger
from genesis.runtime.simclock import SimClock

logger = get_logger("runtime.scheduler")


class SimScheduler:
    def __init__(self, clock: SimClock) -> None:
        self.clock = clock
        self._heap: list[tuple[float, int, str, str]] = []  # (wake_sim_s, seq, agent_id, reason)
        self._seq = itertools.count()
        self._cancelled: set[tuple[str, int]] = set()
        self._pending_seq: dict[str, int] = {}   # agent -> 最新预约 seq(新预约自动作废旧的)
        self._kick = asyncio.Event()
        self.running = False

    # —— 预约/中断 ——
    def schedule(self, agent_id: str, wake_sim_s: float, reason: str = "tick") -> None:
        """预约唤醒;同一 agent 的旧预约自动作废(单一待办,简化心智模型)。"""
        old = self._pending_seq.get(agent_id)
        if old is not None:
            self._cancelled.add((agent_id, old))
        seq = next(self._seq)
        self._pending_seq[agent_id] = seq
        heapq.heappush(self._heap, (wake_sim_s, seq, agent_id, reason))
        self._kick.set()

    def interrupt(self, agent_id: str, reason: str) -> None:
        """高显著度中断:立即唤醒(取代其原有预约)。"""
        self.schedule(agent_id, self.clock.sim_seconds, f"interrupt:{reason}")

    def drop(self, agent_id: str) -> None:
        """agent 退场(死亡):作废其预约。"""
        old = self._pending_seq.pop(agent_id, None)
        if old is not None:
            self._cancelled.add((agent_id, old))

    # —— 主循环 ——
    async def run(self, handler) -> None:
        """handler(agent_id, reason) 为 async 回调;并发触发,互不阻塞。"""
        self.running = True
        while self.running:
            while self._heap and self._heap[0][0] <= self.clock.sim_seconds:
                wake_at, seq, aid, reason = heapq.heappop(self._heap)
                if (aid, seq) in self._cancelled:
                    self._cancelled.discard((aid, seq))
                    continue
                if self._pending_seq.get(aid) == seq:
                    self._pending_seq.pop(aid, None)
                asyncio.create_task(self._safe(handler, aid, reason))
            self._kick.clear()
            delay = self.clock.real_seconds_until(self._heap[0][0]) if self._heap else 0.5
            try:  # 睡到堆顶到点,或被新预约/中断踢醒
                await asyncio.wait_for(self._kick.wait(), timeout=min(max(delay, 0.02), 0.5))
            except asyncio.TimeoutError:
                pass

    async def _safe(self, handler, aid: str, reason: str) -> None:
        try:
            await handler(aid, reason)
        except Exception:  # 单个 agent 异常不拖垮世界
            logger.exception("agent[%s] 唤醒处理失败(%s)", aid, reason)

    def stop(self) -> None:
        self.running = False
        self._kick.set()