"""跟踪 Platform IO 层的出站投递状态。

当前实现基于两组 ``dict + heapq``：
- ``_pending`` 和 ``_pending_expire_heap`` 负责管理待完成的出站记录
- ``_receipts_by_external_id`` 和 ``_receipt_expire_heap`` 负责管理已完成回执索引

这样就不需要在每次读写时全表扫描过期项，而是通过懒清理逐步弹出已经过期
或已经失效的堆节点。
"""

from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple

import heapq
import time

from .types import DeliveryReceipt, RouteKey


@dataclass(slots=True)
class PendingOutboundRecord:
    """表示一条仍在等待完成的出站投递记录。

    Attributes:
        internal_message_id: 正在跟踪的内部 ``SessionMessage.message_id``。
        route_key: 该出站投递开始时使用的路由键。
        driver_id: 负责这次出站投递的驱动 ID。
        created_at: 开始跟踪时记录的单调时钟时间戳。
        expires_at: 该待完成记录预计过期的单调时钟时间戳。
        metadata: 与待完成记录一同保留的额外 Broker 侧元数据。
    """

    internal_message_id: str
    route_key: RouteKey
    driver_id: str
    created_at: float = field(default_factory=time.monotonic)
    expires_at: float = 0.0
    metadata: Dict[str, Any] = field(default_factory=dict)


@dataclass(slots=True)
class StoredDeliveryReceipt:
    """表示一条已完成并暂存的出站回执。

    Attributes:
        receipt: 规范化后的出站投递回执。
        stored_at: 回执被写入索引时记录的单调时钟时间戳。
        expires_at: 该回执索引预计过期的单调时钟时间戳。
    """

    receipt: DeliveryReceipt
    stored_at: float = field(default_factory=time.monotonic)
    expires_at: float = 0.0


class OutboundTracker:
    """统一跟踪出站消息的 pending 状态与最终回执。

    主要用于解决出站消息在发送过程中“状态散落在不同路径里”的问题：
    - 发送开始后，需要在最终回执返回前保留一份 pending 状态
    - 平台返回 ``external_message_id`` 后，需要保留一段时间的回执索引

    当前实现使用 ``dict + heapq`` 做 TTL 管理：
    - ``dict`` 提供 ``O(1)`` 级别的主键查询
    - ``heapq`` 提供按过期时间排序的懒清理能力

    这比“每次 begin/finish/get 都全表扫描”的实现更适合高吞吐出站场景。

    Notes:
        复杂度说明如下，设 ``p`` 为当前有效 pending 数量，``r`` 为当前有效回执数量：

        - ``begin_tracking()``、``finish_tracking()`` 的常见路径时间复杂度接近
          ``O(log p)`` 或 ``O(log r)``
        - ``get_pending()``、``get_receipt_by_external_id()`` 的查询本身是 ``O(1)``
          ，连同懒清理一起看，长期摊还复杂度接近 ``O(log n)``
        - 如果某次调用恰好触发一批过期节点的集中清理，则该次调用的最坏时间复杂度
          可达到 ``O(k log n)``，其中 ``k`` 为本次被弹出的节点数量
        - 空间复杂度为 ``O(p + r)``
    """

    def __init__(self, ttl_seconds: float = 1800.0) -> None:
        """初始化出站跟踪器。

        Args:
            ttl_seconds: 待完成记录与按外部消息 ID 建立的回执索引保留时长，
                单位为秒。

        Raises:
            ValueError: 当 ``ttl_seconds`` 非正数时抛出。
        """
        if ttl_seconds <= 0:
            raise ValueError("ttl_seconds 必须大于 0")

        self._ttl_seconds = ttl_seconds
        self._pending: Dict[Tuple[str, str], PendingOutboundRecord] = {}
        self._pending_expire_heap: List[Tuple[float, str, str]] = []
        self._receipts_by_external_id: Dict[str, StoredDeliveryReceipt] = {}
        self._receipt_expire_heap: List[Tuple[float, str]] = []

    @staticmethod
    def _build_pending_key(internal_message_id: str, driver_id: str) -> Tuple[str, str]:
        """构造单条出站跟踪记录的唯一键。

        Args:
            internal_message_id: 内部消息 ID。
            driver_id: 负责当前投递的驱动 ID。

        Returns:
            Tuple[str, str]: ``(internal_message_id, driver_id)`` 组合键。
        """
        return internal_message_id, driver_id

    def begin_tracking(
        self,
        internal_message_id: str,
        route_key: RouteKey,
        driver_id: str,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> PendingOutboundRecord:
        """开始跟踪一次出站投递。

        Args:
            internal_message_id: 正在投递的内部消息 ID。
            route_key: 这次出站投递选择的路由键。
            driver_id: 负责本次投递的驱动 ID。
            metadata: 可选的额外元数据，会一并保存在待完成记录中。

        Returns:
            PendingOutboundRecord: 新创建的待完成记录。

        Raises:
            ValueError: 当同一个 ``internal_message_id`` 与 ``driver_id`` 组合已经存在
                未完成记录时抛出。
        """
        now = time.monotonic()
        self._cleanup_expired(now)
        pending_key = self._build_pending_key(internal_message_id, driver_id)

        if pending_key in self._pending:
            raise ValueError(f"消息 {internal_message_id} 在驱动 {driver_id} 上已存在未完成的出站跟踪记录")

        expires_at = now + self._ttl_seconds
        record = PendingOutboundRecord(
            internal_message_id=internal_message_id,
            route_key=route_key,
            driver_id=driver_id,
            created_at=now,
            expires_at=expires_at,
            metadata=metadata or {},
        )
        self._pending[pending_key] = record
        heapq.heappush(self._pending_expire_heap, (expires_at, internal_message_id, driver_id))
        return record

    def finish_tracking(self, receipt: DeliveryReceipt) -> Optional[PendingOutboundRecord]:
        """使用最终回执结束一条出站跟踪。

        Args:
            receipt: 规范化后的最终投递回执。

        Returns:
            Optional[PendingOutboundRecord]: 若此前存在待完成记录，则返回该记录。
        """
        now = time.monotonic()
        self._cleanup_expired(now)

        pending_record: Optional[PendingOutboundRecord] = None
        if receipt.driver_id:
            pending_key = self._build_pending_key(receipt.internal_message_id, receipt.driver_id)
            pending_record = self._pending.pop(pending_key, None)
        else:
            matched_records = [
                key
                for key, record in self._pending.items()
                if record.internal_message_id == receipt.internal_message_id
            ]
            if len(matched_records) == 1:
                pending_record = self._pending.pop(matched_records[0], None)

        if receipt.external_message_id:
            expires_at = now + self._ttl_seconds
            self._receipts_by_external_id[receipt.external_message_id] = StoredDeliveryReceipt(
                receipt=receipt,
                stored_at=now,
                expires_at=expires_at,
            )
            heapq.heappush(self._receipt_expire_heap, (expires_at, receipt.external_message_id))
        return pending_record

    def get_pending(
        self,
        internal_message_id: str,
        driver_id: Optional[str] = None,
    ) -> Optional[PendingOutboundRecord]:
        """根据内部消息 ID 查询待完成记录。

        Args:
            internal_message_id: 要查询的内部消息 ID。
            driver_id: 可选的驱动 ID；提供后仅返回该驱动上的待完成记录。

        Returns:
            Optional[PendingOutboundRecord]: 若记录仍存在，则返回对应待完成记录。
        """
        self._cleanup_expired(time.monotonic())

        if driver_id:
            return self._pending.get(self._build_pending_key(internal_message_id, driver_id))

        matched_records = [
            record
            for record in self._pending.values()
            if record.internal_message_id == internal_message_id
        ]
        if len(matched_records) == 1:
            return matched_records[0]
        return None

    def get_receipt_by_external_id(self, external_message_id: str) -> Optional[DeliveryReceipt]:
        """根据外部平台消息 ID 查询已完成回执。

        Args:
            external_message_id: 要查询的平台侧消息 ID。

        Returns:
            Optional[DeliveryReceipt]: 若存在对应回执，则返回该回执。
        """
        self._cleanup_expired(time.monotonic())
        stored_receipt = self._receipts_by_external_id.get(external_message_id)
        return stored_receipt.receipt if stored_receipt else None

    def clear(self) -> None:
        """清空全部待完成记录与已保存回执。"""
        self._pending.clear()
        self._pending_expire_heap.clear()
        self._receipts_by_external_id.clear()
        self._receipt_expire_heap.clear()

    def _cleanup_expired(self, now: float) -> None:
        """清理内存中已经过期的待完成记录与已保存回执。

        Args:
            now: 当前单调时钟时间戳。
        """
        self._cleanup_expired_pending(now)
        self._cleanup_expired_receipts(now)

    def _cleanup_expired_pending(self, now: float) -> None:
        """清理已经过期的待完成记录。

        Args:
            now: 当前单调时钟时间戳。

        Notes:
            堆中可能存在已经失效的旧节点。例如某条记录提前 ``finish`` 后，
            它原本的过期节点仍可能留在堆里。这里会通过和 ``dict`` 中当前记录的
            ``expires_at`` 对比，跳过这类旧节点。
        """
        while self._pending_expire_heap and self._pending_expire_heap[0][0] <= now:
            expires_at, internal_message_id, driver_id = heapq.heappop(self._pending_expire_heap)
            pending_key = self._build_pending_key(internal_message_id, driver_id)
            current_record = self._pending.get(pending_key)
            if current_record is None:
                continue
            if current_record.expires_at != expires_at:
                continue
            self._pending.pop(pending_key, None)

    def _cleanup_expired_receipts(self, now: float) -> None:
        """清理已经过期的回执索引。

        Args:
            now: 当前单调时钟时间戳。

        Notes:
            同一个 ``external_message_id`` 在极端情况下可能被重复写入索引，
            因此这里同样需要通过 ``expires_at`` 和当前 ``dict`` 中的值比对，
            跳过已经失效的旧堆节点。
        """
        while self._receipt_expire_heap and self._receipt_expire_heap[0][0] <= now:
            expires_at, external_message_id = heapq.heappop(self._receipt_expire_heap)
            current_receipt = self._receipts_by_external_id.get(external_message_id)
            if current_receipt is None:
                continue
            if current_receipt.expires_at != expires_at:
                continue
            self._receipts_by_external_id.pop(external_message_id, None)
