from __future__ import annotations

import time
from typing import TYPE_CHECKING

from packaging.version import Version

from limits.typing import Optional, RedisClient, ScriptP, Tuple, Type, Union

from ..util import get_package_data
from .base import MovingWindowSupport, Storage

if TYPE_CHECKING:
    import redis


class RedisInteractor:
    RES_DIR = "resources/redis/lua_scripts"

    SCRIPT_MOVING_WINDOW = get_package_data(f"{RES_DIR}/moving_window.lua")
    SCRIPT_ACQUIRE_MOVING_WINDOW = get_package_data(
        f"{RES_DIR}/acquire_moving_window.lua"
    )
    SCRIPT_CLEAR_KEYS = get_package_data(f"{RES_DIR}/clear_keys.lua")
    SCRIPT_INCR_EXPIRE = get_package_data(f"{RES_DIR}/incr_expire.lua")

    lua_moving_window: ScriptP[Tuple[int, int]]
    lua_acquire_window: ScriptP[bool]

    PREFIX = "LIMITS"

    def prefixed_key(self, key: str) -> str:
        return f"{self.PREFIX}:{key}"

    def get_moving_window(self, key: str, limit: int, expiry: int) -> Tuple[float, int]:
        """
        returns the starting point and the number of entries in the moving
        window

        :param key: rate limit key
        :param expiry: expiry of entry
        :return: (start of window, number of acquired entries)
        """
        key = self.prefixed_key(key)
        timestamp = time.time()
        if window := self.lua_moving_window([key], [timestamp - expiry, limit]):
            return float(window[0]), window[1]

        return timestamp, 0

    def _incr(
        self,
        key: str,
        expiry: int,
        connection: RedisClient,
        elastic_expiry: bool = False,
        amount: int = 1,
    ) -> int:
        """
        increments the counter for a given rate limit key

        :param connection: Redis connection
        :param key: the key to increment
        :param expiry: amount in seconds for the key to expire in
        :param amount: the number to increment by
        """
        key = self.prefixed_key(key)
        value = connection.incrby(key, amount)

        if elastic_expiry or value == amount:
            connection.expire(key, expiry)

        return value

    def _get(self, key: str, connection: RedisClient) -> int:
        """
        :param connection: Redis connection
        :param key: the key to get the counter value for
        """

        key = self.prefixed_key(key)
        return int(connection.get(key) or 0)

    def _clear(self, key: str, connection: RedisClient) -> None:
        """
        :param key: the key to clear rate limits for
        :param connection: Redis connection
        """
        key = self.prefixed_key(key)
        connection.delete(key)

    def _acquire_entry(
        self,
        key: str,
        limit: int,
        expiry: int,
        connection: RedisClient,
        amount: int = 1,
    ) -> bool:
        """
        :param key: rate limit key to acquire an entry in
        :param limit: amount of entries allowed
        :param expiry: expiry of the entry
        :param connection: Redis connection
        :param amount: the number of entries to acquire
        """
        key = self.prefixed_key(key)
        timestamp = time.time()
        acquired = self.lua_acquire_window([key], [timestamp, limit, expiry, amount])

        return bool(acquired)

    def _get_expiry(self, key: str, connection: RedisClient) -> float:
        """
        :param key: the key to get the expiry for
        :param connection: Redis connection
        """

        key = self.prefixed_key(key)
        return max(connection.ttl(key), 0) + time.time()

    def _check(self, connection: RedisClient) -> bool:
        """
        :param connection: Redis connection
        check if storage is healthy
        """
        try:
            return connection.ping()
        except:  # noqa
            return False


class RedisStorage(RedisInteractor, Storage, MovingWindowSupport):
    """
    Rate limit storage with redis as backend.

    Depends on :pypi:`redis`.
    """

    STORAGE_SCHEME = ["redis", "rediss", "redis+unix"]
    """The storage scheme for redis"""

    DEPENDENCIES = {"redis": Version("3.0")}

    def __init__(
        self,
        uri: str,
        connection_pool: Optional[redis.connection.ConnectionPool] = None,
        wrap_exceptions: bool = False,
        **options: Union[float, str, bool],
    ) -> None:
        """
        :param uri: uri of the form ``redis://[:password]@host:port``,
         ``redis://[:password]@host:port/db``,
         ``rediss://[:password]@host:port``, ``redis+unix:///path/to/sock`` etc.
         This uri is passed directly to :func:`redis.from_url` except for the
         case of ``redis+unix://`` where it is replaced with ``unix://``.
        :param connection_pool: if provided, the redis client is initialized with
         the connection pool and any other params passed as :paramref:`options`
        :param wrap_exceptions: Whether to wrap storage exceptions in
         :exc:`limits.errors.StorageError` before raising it.
        :param options: all remaining keyword arguments are passed
         directly to the constructor of :class:`redis.Redis`
        :raise ConfigurationError: when the :pypi:`redis` library is not available
        """
        super().__init__(uri, wrap_exceptions=wrap_exceptions, **options)
        self.dependency = self.dependencies["redis"].module

        uri = uri.replace("redis+unix", "unix")

        if not connection_pool:
            self.storage = self.dependency.from_url(uri, **options)
        else:
            self.storage = self.dependency.Redis(
                connection_pool=connection_pool, **options
            )
        self.initialize_storage(uri)

    @property
    def base_exceptions(
        self,
    ) -> Union[Type[Exception], Tuple[Type[Exception], ...]]:  # pragma: no cover
        return self.dependency.RedisError  # type: ignore[no-any-return]

    def initialize_storage(self, _uri: str) -> None:
        self.lua_moving_window = self.storage.register_script(self.SCRIPT_MOVING_WINDOW)
        self.lua_acquire_window = self.storage.register_script(
            self.SCRIPT_ACQUIRE_MOVING_WINDOW
        )
        self.lua_clear_keys = self.storage.register_script(self.SCRIPT_CLEAR_KEYS)
        self.lua_incr_expire = self.storage.register_script(
            RedisStorage.SCRIPT_INCR_EXPIRE
        )

    def incr(
        self, key: str, expiry: int, elastic_expiry: bool = False, amount: int = 1
    ) -> int:
        """
        increments the counter for a given rate limit key

        :param key: the key to increment
        :param expiry: amount in seconds for the key to expire in
        :param amount: the number to increment by
        """

        if elastic_expiry:
            return super()._incr(key, expiry, self.storage, elastic_expiry, amount)
        else:
            key = self.prefixed_key(key)
            return int(self.lua_incr_expire([key], [expiry, amount]))

    def get(self, key: str) -> int:
        """
        :param key: the key to get the counter value for
        """

        return super()._get(key, self.storage)

    def clear(self, key: str) -> None:
        """
        :param key: the key to clear rate limits for
        """

        return super()._clear(key, self.storage)

    def acquire_entry(self, key: str, limit: int, expiry: int, amount: int = 1) -> bool:
        """
        :param key: rate limit key to acquire an entry in
        :param limit: amount of entries allowed
        :param expiry: expiry of the entry
        :param amount: the number to increment by
        """

        return super()._acquire_entry(key, limit, expiry, self.storage, amount)

    def get_expiry(self, key: str) -> float:
        """
        :param key: the key to get the expiry for
        """

        return super()._get_expiry(key, self.storage)

    def check(self) -> bool:
        """
        check if storage is healthy
        """

        return super()._check(self.storage)

    def reset(self) -> Optional[int]:
        """
        This function calls a Lua Script to delete keys prefixed with
        ``self.PREFIX`` in blocks of 5000.

        .. warning::
           This operation was designed to be fast, but was not tested
           on a large production based system. Be careful with its usage as it
           could be slow on very large data sets.

        """

        prefix = self.prefixed_key("*")
        return int(self.lua_clear_keys([prefix]))
