"""混合确认注册表 — 同进程 Event 唤醒 + 重启后 DB polling 兜底。

Agent 循环在工具需要确认时：
1. 写 PendingConfirmationRecord 到 DB（持久化挂起状态）
2. 调 ``ConfirmationRegistry.wait(run_id)`` — 同进程用 Event 0 延迟唤醒，
   重启后 polling DB 兜底

approve/deny 端点统一调 ``registry.resolve(run_id, approved)``：
- 更新 PendingConfirmationRecord.status + resolved_at
- 同一进程直接 set Event
"""

from __future__ import annotations

import asyncio
import time
from datetime import UTC, datetime

from sqlalchemy import or_, select, update

from app.models.safety import PendingConfirmationRecord


class ConfirmationRegistry:
    """混合确认注册表。

    - 同进程：``asyncio.Event`` 直接唤醒，0 延迟
    - 重启后：Event 不存在 → polling ``PendingConfirmationRecord.status`` 兜底
    - approve/deny 端点统一调用 ``resolve(run_id, approved)``
    """

    def __init__(self) -> None:
        self._events: dict[str, asyncio.Event] = {}

    async def wait(
        self,
        run_id: str,
        *,
        poll_interval: float = 0.5,
        timeout: float = 60,
        confirmation_id: str | None = None,
    ) -> bool:
        """等待确认。

        同一进程的事件直接唤醒；跨进程或重启后 polling 兜底。

        ``confirmation_id`` 可选，传入时会同时注册为该 Event 的别名，
        使得外部以 ``confirmation_id`` 调 ``resolve`` 时也能立即唤醒同进程的 Event
        （无需等 DB polling 兜底）。

        返回：
            True=批准, False=拒绝 ; 超时抛出 TimeoutError
        """
        event = asyncio.Event()
        self._events[run_id] = event
        if confirmation_id:
            self._events[confirmation_id] = event

        try:
            deadline = time.monotonic() + timeout
            while time.monotonic() < deadline:
                # 1) polling DB 获取真实状态（Event 只用于唤醒，不携带批准/拒绝值）
                from app.db.session import AsyncSessionLocal

                async with AsyncSessionLocal() as session:
                    if confirmation_id is None:
                        stmt = select(PendingConfirmationRecord).where(
                            or_(
                                PendingConfirmationRecord.id == run_id,
                                PendingConfirmationRecord.run_id == run_id,
                            )
                        )
                    else:
                        stmt = select(PendingConfirmationRecord).where(
                            PendingConfirmationRecord.id == confirmation_id,
                            PendingConfirmationRecord.run_id == run_id,
                        )
                    result = await session.execute(stmt)
                    record = result.scalar_one_or_none()
                    if record is not None and record.status != "pending":
                        return record.status == "approved"

                # 2) 等一个间隔，同时监听 Event（Event set 只表示「有变化」，实际值靠 DB）
                try:
                    await asyncio.wait_for(
                        asyncio.shield(event.wait()),
                        timeout=poll_interval,
                    )
                    # Event 被唤醒 → 回到循环顶部重新查 DB
                    continue
                except TimeoutError:
                    continue

            # 超时
            raise TimeoutError(f"确认超时（等待 {timeout:.0f}s 未裁决）")
        finally:
            self._events.pop(run_id, None)
            if confirmation_id:
                self._events.pop(confirmation_id, None)

    async def resolve(
        self,
        run_id: str,
        approved: bool,
        *,
        confirmation_id: str | None = None,
    ) -> bool:
        """精确裁决一个待确认工具调用，并唤醒对应等待者。"""
        from app.db.session import AsyncSessionLocal

        target_id = confirmation_id or run_id
        async with AsyncSessionLocal() as session:
            status = "approved" if approved else "denied"
            stmt = update(PendingConfirmationRecord).where(
                PendingConfirmationRecord.status == "pending"
            )
            if confirmation_id is None:
                stmt = stmt.where(
                    or_(
                        PendingConfirmationRecord.id == run_id,
                        PendingConfirmationRecord.run_id == run_id,
                    )
                )
            else:
                stmt = stmt.where(
                    PendingConfirmationRecord.id == confirmation_id,
                    PendingConfirmationRecord.run_id == run_id,
                )
            result = await session.execute(
                stmt.values(status=status, resolved_at=datetime.now(UTC))
            )
            await session.commit()

        if result.rowcount == 0:
            return False

        event = self._events.get(target_id) or self._events.get(run_id)
        if event is not None:
            event.set()
        return True

    def has(self, run_id: str) -> bool:
        """是否仍在等待（同进程检查）。"""
        event = self._events.get(run_id)
        return event is not None and not event.is_set()


__all__ = ["ConfirmationRegistry"]