"""API Key 池管理。"""

from __future__ import annotations

import hashlib
from collections.abc import Iterable

from django.core.cache import cache

_PREFERRED_KEY_CACHE_TTL_SECONDS = 30 * 24 * 60 * 60
_AUTH_BLOCK_TTL_SECONDS = 60 * 60
_RATE_LIMIT_BLOCK_TTL_SECONDS = 2 * 60


class McpApiKeyPool:
    """为同一 provider 管理多个 API Key 的优先级与短期熔断。"""

    def __init__(self, *, provider_name: str, api_keys: Iterable[str]) -> None:
        self._provider_name = str(provider_name or "").strip().lower()
        self._api_keys = self._normalize_api_keys(api_keys)

    @property
    def size(self) -> int:
        return len(self._api_keys)

    def ordered_keys(self) -> list[str]:
        if len(self._api_keys) <= 1:
            return list(self._api_keys)

        preferred = str(cache.get(self._preferred_cache_key()) or "").strip()
        available = [key for key in self._api_keys if not self._is_blocked(key)]
        blocked = [key for key in self._api_keys if self._is_blocked(key)]
        ordered_available = self._order_with_preferred(available, preferred)
        ordered_blocked = self._order_with_preferred(blocked, preferred)
        return ordered_available + ordered_blocked

    def mark_success(self, api_key: str) -> None:
        fingerprint = self.fingerprint(api_key)
        if not fingerprint:
            return
        cache.set(self._preferred_cache_key(), fingerprint, timeout=_PREFERRED_KEY_CACHE_TTL_SECONDS)
        cache.delete(self._block_cache_key(api_key))

    def mark_auth_failed(self, api_key: str) -> None:
        self._block(api_key, ttl_seconds=_AUTH_BLOCK_TTL_SECONDS)

    def mark_rate_limited(self, api_key: str) -> None:
        self._block(api_key, ttl_seconds=_RATE_LIMIT_BLOCK_TTL_SECONDS)

    def fingerprint(self, api_key: str) -> str:
        normalized = str(api_key or "").strip()
        if not normalized:
            return ""
        digest = hashlib.sha256(normalized.encode("utf-8")).hexdigest()
        return digest[:16]

    @staticmethod
    def _normalize_api_keys(api_keys: Iterable[str]) -> tuple[str, ...]:
        normalized: list[str] = []
        seen: set[str] = set()
        for item in api_keys:
            value = str(item or "").strip()
            if not value or value in seen:
                continue
            normalized.append(value)
            seen.add(value)
        return tuple(normalized)

    def _block(self, api_key: str, *, ttl_seconds: int) -> None:
        if ttl_seconds <= 0:
            return
        cache.set(self._block_cache_key(api_key), True, timeout=ttl_seconds)

    def _is_blocked(self, api_key: str) -> bool:
        return bool(cache.get(self._block_cache_key(api_key)))

    def _preferred_cache_key(self) -> str:
        return f"enterprise_data:api_key_pool:{self._provider_name}:preferred"

    def _block_cache_key(self, api_key: str) -> str:
        return f"enterprise_data:api_key_pool:{self._provider_name}:blocked:{self.fingerprint(api_key)}"

    def _order_with_preferred(self, keys: list[str], preferred_fingerprint: str) -> list[str]:
        if len(keys) <= 1 or not preferred_fingerprint:
            return list(keys)
        for index, key in enumerate(keys):
            if self.fingerprint(key) != preferred_fingerprint:
                continue
            return [keys[index], *keys[:index], *keys[index + 1 :]]
        return list(keys)