from __future__ import annotations

import logging
import threading
from typing import Any

_RELAXED_DEFAULT_LIMIT = 1_000_000

logger = logging.getLogger(__name__)


class BaseConcurrencyLimiter:
    def acquire(self, org_id: str) -> bool:
        raise NotImplementedError

    def release(self, org_id: str) -> None:
        raise NotImplementedError

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


class InMemoryConcurrencyLimiter(BaseConcurrencyLimiter):
    def __init__(self, default_limit: int = _RELAXED_DEFAULT_LIMIT):
        self._limits: dict[str, int] = {}
        self._counts: dict[str, int] = {}
        self._lock = threading.Lock()
        self._default = int(default_limit)

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

    def acquire(self, org_id: str) -> bool:
        with self._lock:
            limit = int(self._limits.get(org_id, self._default))
            cur = int(self._counts.get(org_id, 0))
            if limit <= 0:
                return False
            if cur >= limit:
                return False
            self._counts[org_id] = cur + 1
            return True

    def release(self, org_id: str) -> None:
        with self._lock:
            cur = int(self._counts.get(org_id, 0))
            self._counts[org_id] = max(0, cur - 1)


_ACQUIRE_SCRIPT = """\
local cur = redis.call('INCR', KEYS[1])
redis.call('EXPIRE', KEYS[1], tonumber(ARGV[2]))
if cur > tonumber(ARGV[1]) then
    redis.call('DECR', KEYS[1])
    return 0
end
return 1
"""

_RELEASE_SCRIPT = """\
local val = redis.call('GET', KEYS[1])
if val == false or tonumber(val) <= 0 then
    redis.call('DEL', KEYS[1])
    return 0
end
return redis.call('DECR', KEYS[1])
"""


class RedisConcurrencyLimiter(BaseConcurrencyLimiter):
    """Redis-backed concurrency limiter using atomic Lua scripts.

    Uses key pattern: ccy:{org_id} with TTL refreshed on every acquire to
    prevent counter drift from long-running requests.

    Both acquire and release are atomic: the Lua script performs the INCR,
    limit check, and rollback in a single Redis operation, eliminating the
    race window where concurrent requests could exceed the limit.
    """

    def __init__(
        self,
        client: Any,
        default_limit: int = _RELAXED_DEFAULT_LIMIT,
        ttl_seconds: int = 300,
    ) -> None:
        self.client = client
        self.default = int(default_limit)
        self.ttl = int(ttl_seconds)
        self._limits: dict[str, int] = {}
        self._acquire_sha: str | None = None
        self._release_sha: str | None = None

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

    def _limit(self, org_id: str) -> int:
        return int(self._limits.get(org_id, self.default))

    def acquire(self, org_id: str) -> bool:
        key = f"ccy:{org_id}"
        limit = self._limit(org_id)
        try:
            result = self._eval_script(
                _ACQUIRE_SCRIPT, self._acquire_sha,
                keys=[key], args=[str(limit), str(self.ttl)],
            )
            if self._acquire_sha is None:
                self._acquire_sha = self._try_cache_script(_ACQUIRE_SCRIPT)
            return int(result) == 1
        except Exception as exc:
            logger.error(
                "Redis concurrency limiter acquire failed",
                exc_info=exc, extra={"org_id": org_id},
            )
            return False

    def _try_cache_script(self, script: str) -> str | None:
        """Attempt to cache a script SHA for future evalsha calls."""
        try:
            return self.client.script_load(script)
        except Exception:
            return None

    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 release(self, org_id: str) -> None:
        key = f"ccy:{org_id}"
        try:
            self._eval_script(
                _RELEASE_SCRIPT, self._release_sha,
                keys=[key], args=[],
            )
            if self._release_sha is None:
                self._release_sha = self._try_cache_script(_RELEASE_SCRIPT)
        except Exception as exc:
            # Final fallback: use simple DECR but clamp via GET check first.
            try:
                val = self.client.get(key)
                if val is not None and int(val) > 0:
                    self.client.decr(key)
            except Exception as inner_exc:
                logger.error(
                    "Redis concurrency release fallback also failed — "
                    "slot may be permanently leaked for org_id=%s",
                    org_id,
                    exc_info=inner_exc,
                )
            logger.warning(
                "Redis concurrency limiter release used fallback",
                exc_info=exc,
                extra={"org_id": org_id},
            )
