"""内存缓存与限流计数，减少对数据库的依赖。

key 缓存：首次校验后把 key 的校验信息与配额缓存在内存，
后续请求命中缓存即不查库；后台改 key 时清缓存。
RPM：用内存滑动窗口计数，不再查 request_logs 表。
"""
from __future__ import annotations

import asyncio
import time
from dataclasses import dataclass
from typing import Dict, List, Set

import hashlib


@dataclass
class CachedKey:
    """key 的缓存校验信息（不含敏感的完整 key）。"""
    id: int
    name: str
    prefix: str
    key_hash: str
    is_active: bool
    expires_at: float | None  # epoch 秒
    allowed_models: list[str] | None
    rpm_limit: int | None
    tpm_limit: int | None
    token_quota: int | None
    tokens_used: int
    # 负载快照时间，用于判断是否需要回库刷新 tokens_used
    loaded_at: float


class KeyCache:
    """按 prefix 索引的 key 内存缓存。"""

    def __init__(self) -> None:
        self._by_prefix: Dict[str, CachedKey] = {}
        self._by_id: Dict[int, str] = {}  # id -> prefix，用于按 id 失效
        self._lock = asyncio.Lock()

    def get(self, prefix: str) -> CachedKey | None:
        return self._by_prefix.get(prefix)

    async def put(self, k: CachedKey) -> None:
        async with self._lock:
            self._by_prefix[k.prefix] = k
            self._by_id[k.id] = k.prefix

    async def invalidate_id(self, key_id: int) -> None:
        """按 id 失效（后台改/删 key 时调用）。"""
        async with self._lock:
            prefix = self._by_id.pop(key_id, None)
            if prefix:
                self._by_prefix.pop(prefix, None)

    async def invalidate_prefix(self, prefix: str) -> None:
        async with self._lock:
            ck = self._by_prefix.pop(prefix, None)
            if ck:
                self._by_id.pop(ck.id, None)

    async def clear(self) -> None:
        async with self._lock:
            self._by_prefix.clear()
            self._by_id.clear()

    async def add_usage(self, key_id: int, tokens: int) -> None:
        """成功转发后把用量累加到内存缓存（不回库），保持配额判断近实时。"""
        async with self._lock:
            prefix = self._by_id.get(key_id)
            if prefix:
                ck = self._by_prefix.get(prefix)
                if ck:
                    ck.tokens_used += tokens


class RPMCounter:
    """内存滑动窗口 RPM 计数，按 key_id 维度。"""

    def __init__(self, window: float = 60.0) -> None:
        self._window = window
        self._hits: Dict[int, List[float]] = {}
        self._lock = asyncio.Lock()

    async def hit_and_check(self, key_id: int, limit: int | None) -> bool:
        """记录一次命中并返回是否在限额内（True=放行，False=超限）。"""
        if not limit:
            return True
        now = time.time()
        cutoff = now - self._window
        async with self._lock:
            hits = self._hits.get(key_id)
            if hits is None:
                hits = []
                self._hits[key_id] = hits
            # 清掉过期命中
            hits[:] = [t for t in hits if t > cutoff]
            if len(hits) >= limit:
                return False
            hits.append(now)
            return True


key_cache = KeyCache()
rpm_counter = RPMCounter()


def verify_key_hash(raw: str, digest: str) -> bool:
    return hashlib.sha256(raw.encode()).hexdigest() == digest


async def cached_key_from(raw: str, db_loader) -> CachedKey | None:
    """先查缓存，未命中则调 db_loader(prefix) 回库加载并缓存。"""
    prefix = raw[:20]
    ck = key_cache.get(prefix)
    if ck is not None and verify_key_hash(raw, ck.key_hash):
        return ck
    # 未命中或哈希不匹配，回库
    db_key = await db_loader(prefix)
    if db_key is None:
        return None
    exp = db_key.expires_at.timestamp() if db_key.expires_at else None
    ck = CachedKey(
        id=db_key.id, name=db_key.name, prefix=db_key.prefix, key_hash=db_key.key_hash,
        is_active=db_key.is_active, expires_at=exp, allowed_models=db_key.allowed_models,
        rpm_limit=db_key.rpm_limit, tpm_limit=db_key.tpm_limit,
        token_quota=db_key.token_quota, tokens_used=db_key.tokens_used or 0,
        loaded_at=time.time(),
    )
    await key_cache.put(ck)
    return ck