"""Usage Tracker - 使用量追踪"""

from dataclasses import dataclass, field
from datetime import datetime, timezone

from agentkit.llm.protocol import TokenUsage


@dataclass
class UsageRecord:
    """使用量记录"""

    agent_name: str
    model: str
    prompt_tokens: int
    completion_tokens: int
    total_tokens: int
    cost: float
    latency_ms: float
    timestamp: datetime = field(default_factory=lambda: datetime.now(timezone.utc))


@dataclass
class UsageSummary:
    """使用量汇总"""

    total_tokens: int = 0
    total_cost: float = 0.0
    by_model: dict[str, dict[str, int | float]] = field(default_factory=dict)
    records: list[UsageRecord] = field(default_factory=list)


class UsageTracker:
    """使用量追踪器"""

    MAX_RECORDS = 10000  # 最大记录数，防止内存无限增长

    def __init__(self) -> None:
        self._records: list[UsageRecord] = []

    def record(
        self,
        agent_name: str,
        model: str,
        usage: TokenUsage,
        cost: float,
        latency_ms: float,
    ) -> None:
        """记录一次使用"""
        rec = UsageRecord(
            agent_name=agent_name,
            model=model,
            prompt_tokens=usage.prompt_tokens,
            completion_tokens=usage.completion_tokens,
            total_tokens=usage.total_tokens,
            cost=cost,
            latency_ms=latency_ms,
        )
        self._records.append(rec)
        # 超过上限时删除最早的记录
        if len(self._records) > self.MAX_RECORDS:
            self._records = self._records[-self.MAX_RECORDS:]

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

        if agent_name is not None:
            filtered = [r for r in filtered if r.agent_name == agent_name]
        if start_time is not None:
            filtered = [r for r in filtered if r.timestamp >= start_time]
        if end_time is not None:
            filtered = [r for r in filtered if r.timestamp <= end_time]

        if not filtered:
            return UsageSummary()

        total_tokens = sum(r.total_tokens for r in filtered)
        total_cost = sum(r.cost for r in filtered)

        by_model: dict[str, dict[str, int | float]] = {}
        for r in filtered:
            if r.model not in by_model:
                by_model[r.model] = {"total_tokens": 0, "total_cost": 0.0, "count": 0}
            by_model[r.model]["total_tokens"] += r.total_tokens
            by_model[r.model]["total_cost"] += r.cost
            by_model[r.model]["count"] += 1

        return UsageSummary(
            total_tokens=total_tokens,
            total_cost=total_cost,
            by_model=by_model,
            records=filtered,
        )
