from __future__ import annotations

import logging
import threading
import time
from dataclasses import dataclass
from typing import Any

from .settings import get_gateway_settings

logger = logging.getLogger(__name__)


class BaseTPMLimiter:
    """Abstract tokens-per-minute limiter interface.

    Implementations should be safe for multi-request access. For distributed
    deployments use a Redis-backed implementation; this in-memory version is
    suitable for single-process or test environments.
    """

    def consume(self, org_id: str, tokens: int) -> bool:  # noqa: D401
        """Consume `tokens` from the org's minute budget."""
        raise NotImplementedError

    def consume_input(self, org_id: str, tokens: int) -> bool:
        return self.consume(org_id, tokens)

    def consume_output(self, org_id: str, tokens: int) -> bool:
        return self.consume(org_id, tokens)

    def set_limit(self, org_id: str, tokens_per_minute: int) -> None:
        raise NotImplementedError

    def set_input_limit(self, org_id: str, tokens_per_minute: int) -> None:
        self.set_limit(org_id, tokens_per_minute)

    def set_output_limit(self, org_id: str, tokens_per_minute: int) -> None:
        self.set_limit(org_id, tokens_per_minute)

    def get_limit(self, org_id: str) -> int | None:
        raise NotImplementedError

    def reconcile(self, org_id: str, precharged: int, actual: int) -> None:
        """Refund the difference between pre-charged and actual token usage.

        Called after a request completes to correct the pre-charge estimate.
        Subclasses should override with their storage-specific refund logic.

        Args:
            org_id: Organization identifier.
            precharged: Number of tokens pre-charged before the request.
            actual: Actual tokens consumed by the request.
        """
        diff = max(0, precharged - actual)
        if diff > 0:
            self._refund(org_id, diff)

    def _refund(self, org_id: str, tokens: int) -> None:
        """Refund tokens to an org's budget. Override in subclasses."""


@dataclass(slots=True)
class _Window:
    buckets: list[int]
    bucket_seconds: list[int]


class InMemoryTPMLimiter(BaseTPMLimiter):
    """Sliding window TPM limiter using a fixed 60-second ring per org."""

    def __init__(self, default_limit: int | None = None, window_seconds: int = 60):
        self.window_seconds = max(1, window_seconds)
        settings = get_gateway_settings()
        effective_default = default_limit
        if effective_default is None:
            configured = settings.limits.default_tpm
            effective_default = None if configured is None else int(configured)
        self.default_limit: int | None = (
            int(effective_default) if effective_default is not None else None
        )
        self._lock = threading.Lock()
        self._limits: dict[str, int] = {}
        self._windows: dict[str, _Window] = {}

    def _sec(self) -> int:
        return int(time.time())

    def set_limit(self, org_id: str, tokens_per_minute: int) -> None:
        with self._lock:
            self._limits[org_id] = max(0, int(tokens_per_minute))

    def get_limit(self, org_id: str) -> int | None:
        with self._lock:
            explicit = self._limits.get(org_id)
            if explicit is not None:
                return int(explicit)
            if self.default_limit is None:
                return None
            return int(self.default_limit)

    def _get_window_locked(self, org_id: str) -> _Window:
        win = self._windows.get(org_id)
        if win is not None:
            return win
        win = _Window(buckets=[0] * self.window_seconds, bucket_seconds=[0] * self.window_seconds)
        self._windows[org_id] = win
        return win

    def consume(self, org_id: str, tokens: int) -> bool:
        tokens = max(0, int(tokens))
        if tokens == 0:
            return True
        limit = self.get_limit(org_id)
        if limit is None:
            return True
        if limit <= 0:
            return False
        now = self._sec()
        with self._lock:
            win = self._get_window_locked(org_id)
            idx = now % self.window_seconds
            if win.bucket_seconds[idx] != now:
                win.bucket_seconds[idx] = now
                win.buckets[idx] = 0

            used = sum(
                count
                for count, bucket_sec in zip(win.buckets, win.bucket_seconds, strict=False)
                if now - bucket_sec < self.window_seconds
            )
            if used + tokens > limit:
                return False

            win.buckets[idx] += tokens
            return True

    def _refund(self, org_id: str, tokens: int) -> None:
        now = self._sec()
        with self._lock:
            win = self._get_window_locked(org_id)
            idx = now % self.window_seconds
            if win.bucket_seconds[idx] != now:
                win.bucket_seconds[idx] = now
                win.buckets[idx] = 0
            win.buckets[idx] = max(0, win.buckets[idx] - tokens)


class DualTPMLimiter(BaseTPMLimiter):
    """Dual-channel TPM limiter (input/output) for providers like Anthropic."""

    def __init__(self, input_limiter: BaseTPMLimiter, output_limiter: BaseTPMLimiter):
        self._in = input_limiter
        self._out = output_limiter

    def set_input_limit(self, org_id: str, tokens_per_minute: int) -> None:
        self._in.set_limit(org_id, tokens_per_minute)

    def set_output_limit(self, org_id: str, tokens_per_minute: int) -> None:
        self._out.set_limit(org_id, tokens_per_minute)

    def consume_input(self, org_id: str, tokens: int) -> bool:
        return self._in.consume(org_id, tokens)

    def consume_output(self, org_id: str, tokens: int) -> bool:
        return self._out.consume(org_id, tokens)

    def consume(self, org_id: str, tokens: int) -> bool:
        return self.consume_output(org_id, tokens)

    def set_limit(self, org_id: str, tokens_per_minute: int) -> None:
        self.set_input_limit(org_id, tokens_per_minute)
        self.set_output_limit(org_id, tokens_per_minute)

    def get_limit(self, org_id: str) -> int | None:
        return self._out.get_limit(org_id)

    def reconcile(self, org_id: str, precharged: int, actual: int) -> None:
        self._in.reconcile(org_id, precharged, actual)

    def _refund(self, org_id: str, tokens: int) -> None:
        self._in._refund(org_id, tokens)


_TPM_CONSUME_SCRIPT = """\
local cur = redis.call('INCRBY', KEYS[1], tonumber(ARGV[1]))
redis.call('EXPIRE', KEYS[1], 120)
if cur > tonumber(ARGV[2]) then
    redis.call('DECRBY', KEYS[1], tonumber(ARGV[1]))
    return 0
end
return 1
"""

_TPM_REFUND_SCRIPT = """\
local val = redis.call('GET', KEYS[1])
if val == false then return 0 end
local cur = tonumber(val)
local refund = math.min(tonumber(ARGV[1]), cur)
if refund > 0 then
    redis.call('DECRBY', KEYS[1], refund)
end
return refund
"""


class RedisTPMLimiter(BaseTPMLimiter):
    """Redis-backed fixed-window TPM limiter (minute granularity).

    Uses atomic Lua scripts for consume and refund to eliminate TOCTOU
    races under concurrent load.
    """

    def __init__(
        self,
        client: Any,
        prefix: str = "tpm",
        default_limit: int | None = None,
    ) -> None:
        self.client = client
        self.prefix = prefix
        self.default_limit: int | None = int(default_limit) if default_limit is not None else None
        self._limits: dict[str, int] = {}
        self._consume_sha: str | None = None
        self._refund_sha: str | None = None

    def _minute_key(self, org_id: str) -> str:
        minute = int(time.time() // 60)
        return f"{self.prefix}:{org_id}:{minute}"

    def set_limit(self, org_id: str, tokens_per_minute: int) -> None:
        self._limits[org_id] = max(0, int(tokens_per_minute))

    def get_limit(self, org_id: str) -> int | None:
        explicit = self._limits.get(org_id)
        if explicit is not None:
            return int(explicit)
        if self.default_limit is None:
            return None
        return int(self.default_limit)

    def _eval_script(
        self,
        script: str,
        cached_sha: str | None,
        *,
        keys: list[str],
        args: list[str],
    ) -> Any:
        """Execute a Lua script, preferring cached SHA when available."""
        num_keys = len(keys)
        if cached_sha:
            try:
                return self.client.evalsha(cached_sha, num_keys, *keys, *args)
            except Exception:
                logger.debug(
                    "Redis evalsha failed, falling back to eval",
                    exc_info=True,
                )
        return self.client.eval(script, num_keys, *keys, *args)

    def _try_cache_script(self, script: str) -> str | None:
        try:
            return self.client.script_load(script)
        except Exception:
            logger.warning(
                "Redis script_load failed — Lua scripts will use eval "
                "(slower) on every call",
                exc_info=True,
            )
            return None

    def consume(self, org_id: str, tokens: int) -> bool:
        tokens = max(0, int(tokens))
        if tokens == 0:
            return True
        limit = self.get_limit(org_id)
        if limit is None:
            return True
        if limit <= 0:
            return False
        key = self._minute_key(org_id)
        try:
            result = self._eval_script(
                _TPM_CONSUME_SCRIPT, self._consume_sha,
                keys=[key], args=[str(tokens), str(limit)],
            )
            if self._consume_sha is None:
                self._consume_sha = self._try_cache_script(_TPM_CONSUME_SCRIPT)
            return int(result) == 1
        except Exception as exc:
            logger.error(
                "Redis TPM limiter consume failed",
                exc_info=exc, extra={"org_id": org_id},
            )
            return False

    def _refund(self, org_id: str, tokens: int) -> None:
        key = self._minute_key(org_id)
        try:
            self._eval_script(
                _TPM_REFUND_SCRIPT, self._refund_sha,
                keys=[key], args=[str(tokens)],
            )
            if self._refund_sha is None:
                self._refund_sha = self._try_cache_script(_TPM_REFUND_SCRIPT)
        except Exception as exc:
            logger.warning(
                "Redis TPM limiter refund failed",
                exc_info=exc, extra={"org_id": org_id, "tokens": tokens},
            )
