"""Usage Tracker - 使用量追踪（委托给 UsageStore）"""

from datetime import datetime

from agentkit.llm.protocol import TokenUsage
from agentkit.llm.providers.usage_store import (
    InMemoryUsageStore,
    UsageStore,
    UsageSummary,
)


class UsageTracker:
    """使用量追踪器 — 委托给可插拔的 UsageStore"""

    def __init__(self, store: UsageStore | None = None) -> None:
        self._store: UsageStore = store or InMemoryUsageStore()

    def record(
        self,
        agent_name: str,
        model: str,
        usage: TokenUsage,
        cost: float,
        latency_ms: float,
        user_id: str | None = None,
        department_id: str | None = None,
    ) -> None:
        """记录一次使用（sync — 在 async 上下文中可能阻塞事件循环）"""
        self._store.record(
            agent_name,
            model,
            usage,
            cost,
            latency_ms,
            user_id=user_id,
            department_id=department_id,
        )

    async def record_async(
        self,
        agent_name: str,
        model: str,
        usage: TokenUsage,
        cost: float,
        latency_ms: float,
        user_id: str | None = None,
        department_id: str | None = None,
    ) -> None:
        """记录一次使用（async — 不阻塞事件循环，KTD-6）。

        如果底层 store 提供了 ``record_async`` 方法则调用它；否则
        回退到同步 ``record``（适用于 InMemoryUsageStore 等无 I/O 的 store）。
        """
        record_async = getattr(self._store, "record_async", None)
        if record_async is not None:
            await record_async(
                agent_name,
                model,
                usage,
                cost,
                latency_ms,
                user_id=user_id,
                department_id=department_id,
            )
        else:
            # Fallback: sync store without async support.
            self._store.record(
                agent_name,
                model,
                usage,
                cost,
                latency_ms,
                user_id=user_id,
                department_id=department_id,
            )

    def get_usage(
        self,
        agent_name: str | None = None,
        start_time: datetime | None = None,
        end_time: datetime | None = None,
        user_id: str | None = None,
        department_id: str | None = None,
    ) -> UsageSummary:
        """查询使用量汇总"""
        return self._store.get_usage(
            agent_name=agent_name,
            start_time=start_time,
            end_time=end_time,
            user_id=user_id,
            department_id=department_id,
        )

    @property
    def store(self) -> UsageStore:
        """Expose the underlying store for service-layer queries."""
        return self._store
