from __future__ import annotations

from dataclasses import dataclass, field
from datetime import UTC, datetime, timedelta
from typing import Any

_QUOTA_WARNING_PCT = 80.0
_QUOTA_EXCEEDED_PCT = 100.0

_PERIOD_CUTOFF_DAYS = {"day": 0, "week": 7, "month": 30}


@dataclass
class UsageRecord:
    id: str
    user_id: str
    brand_id: str
    engine_type: str
    query: str
    input_tokens: int
    output_tokens: int
    cost: float
    timestamp: datetime
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass
class UsageSummary:
    period: str
    start_date: str
    end_date: str
    total_queries: int
    total_input_tokens: int
    total_output_tokens: int
    total_cost: float
    by_engine: dict[str, dict[str, Any]]
    by_day: dict[str, dict[str, Any]]


def _aggregate_by_engine(records: list[UsageRecord]) -> dict[str, dict[str, Any]]:
    result: dict[str, dict[str, Any]] = {}
    for r in records:
        bucket = result.setdefault(r.engine_type, {"queries": 0, "input_tokens": 0, "output_tokens": 0, "cost": 0.0})
        bucket["queries"] += 1
        bucket["input_tokens"] += r.input_tokens
        bucket["output_tokens"] += r.output_tokens
        bucket["cost"] += r.cost
    return result


def _compute_cutoff(period: str, now: datetime) -> datetime:
    days = _PERIOD_CUTOFF_DAYS.get(period, 30)
    if days == 0:
        return now.replace(hour=0, minute=0, second=0, microsecond=0)
    return now - timedelta(days=days)


def _quota_status(usage_pct: float) -> str:
    if usage_pct >= _QUOTA_EXCEEDED_PCT:
        return "exceeded"
    if usage_pct >= _QUOTA_WARNING_PCT:
        return "warning"
    return "ok"


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

    def record(
        self,
        user_id: str,
        brand_id: str,
        engine_type: str,
        query: str,
        input_tokens: int,
        output_tokens: int,
        cost: float,
        metadata: dict | None = None,
    ) -> UsageRecord:
        rec = UsageRecord(
            id=f"usage_{len(self._records) + 1}",
            user_id=user_id,
            brand_id=brand_id,
            engine_type=engine_type,
            query=query,
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            cost=cost,
            timestamp=datetime.now(UTC),
            metadata=metadata or {},
        )
        self._records.append(rec)
        return rec

    def get_summary(
        self,
        user_id: str | None = None,
        period: str = "month",
        brand_id: str | None = None,
    ) -> UsageSummary:
        now = datetime.now(UTC)
        filtered = list(self._records)

        if user_id:
            filtered = [r for r in filtered if r.user_id == user_id]
        if brand_id:
            filtered = [r for r in filtered if r.brand_id == brand_id]

        cutoff = _compute_cutoff(period, now)
        filtered = [r for r in filtered if r.timestamp >= cutoff]

        return UsageSummary(
            period=period,
            start_date=cutoff.isoformat(),
            end_date=now.isoformat(),
            total_queries=len(filtered),
            total_input_tokens=sum(r.input_tokens for r in filtered),
            total_output_tokens=sum(r.output_tokens for r in filtered),
            total_cost=round(sum(r.cost for r in filtered), 4),
            by_engine=_aggregate_by_engine(filtered),
            by_day={},
        )

    def check_quota(self, user_id: str, monthly_limit: float = 100.0) -> dict:
        summary = self.get_summary(user_id=user_id, period="month")
        usage_pct = (summary.total_cost / monthly_limit * 100) if monthly_limit > 0 else 0
        return {
            "used": summary.total_cost,
            "limit": monthly_limit,
            "usage_percentage": round(usage_pct, 1),
            "status": _quota_status(usage_pct),
        }