from dataclasses import dataclass
from datetime import datetime, timedelta
from typing import Awaitable, Callable, List, Optional

from quart import current_app, Quart, request, Response
from quart.exceptions import TooManyRequests

from .store import MemoryStore, RateLimiterStoreABC

QUART_RATE_LIMITER_ATTRIBUTE = "_quart_rate_limiter_limits"

KeyCallable = Callable[[], Awaitable[str]]


class RateLimitExceeded(TooManyRequests):
    """A 429 Rate limit exceeded error.

    Arguments:
        retry_after: Seconds left till the remaining resets to the limit.
    """

    def __init__(self, retry_after: int) -> None:
        super().__init__()
        self.retry_after = retry_after

    def get_headers(self) -> dict:
        return {"Retry-After": str(self.retry_after)}


@dataclass
class RateLimit:
    count: int
    period: timedelta
    key_function: Optional[KeyCallable]

    @property
    def inverse(self) -> float:
        return self.period.total_seconds() / self.count

    @property
    def key(self) -> str:
        return f"{self.count}-{self.period.total_seconds()}"


def rate_limit(
    limit: int, period: timedelta, key_function: Optional[KeyCallable] = None
) -> Callable:
    """A decorator to add a rate limit marker to the route.

    This should be used to wrap a route handler (or view function) to
    apply a rate limit to requests to that route. Note that it is
    important that this decorator be wrapped by the route decorator
    and not vice, versa, as below.

    .. code-block:: python

        @app.route('/')
        @rate_limit(10, timedelta(seconds=10))
        async def index():
            ...

    Arguments:
        limit: The maximum number of requests to serve within a
            period.
        period: The duration over which the number of requests must
            not exceed the *limit*.
        key_function: A coroutine function that returns a unique key
            to identify the user agent.

    .. code-block:: python

        async def example_key_function() -> str:
            return request.remote_addr

    """

    def decorator(func: Callable) -> Callable:
        rate_limits = getattr(func, QUART_RATE_LIMITER_ATTRIBUTE, [])
        rate_limits.append(RateLimit(limit, period, key_function))
        setattr(func, QUART_RATE_LIMITER_ATTRIBUTE, rate_limits)
        return func

    return decorator


async def remote_addr_key() -> str:
    return request.remote_addr


class RateLimiter:
    """A Rate limiter instance.

    This can be used to initialise Rate Limiting for a given app,
    either directly,

    .. code-block:: python

        app = Quart(__name__)
        rate_limiter = RateLimiter(app)

    or via the factory pattern,

    .. code-block:: python

        rate_limiter = RateLimiter()

        def create_app():
            app = Quart(__name__)
            rate_limiter.init_app(app)
            return app

    The limiter itself can be customised using the following
    arguments,

    Arguments:
        key_function: A coroutine function that returns a unique key
            to identify the user agent.
        store: The store that contains the theoretical arrival times by
            key.
    """

    def __init__(
        self,
        app: Optional[Quart] = None,
        key_function: KeyCallable = remote_addr_key,
        store: Optional[RateLimiterStoreABC] = None,
    ) -> None:
        self.key_function = key_function
        self.store: RateLimiterStoreABC
        if store is None:
            self.store = MemoryStore()
        else:
            self.store = store

        if app is not None:
            self.init_app(app)

    def init_app(self, app: Quart) -> None:
        app.before_request(self._before_request)
        app.after_request(self._after_request)
        app.before_serving(self._before_serving)
        app.after_serving(self._after_serving)

    async def _before_serving(self) -> None:
        await self.store.before_serving()

    async def _after_serving(self) -> None:
        await self.store.after_serving()

    async def _before_request(self) -> None:
        endpoint = request.endpoint
        view_func = current_app.view_functions.get(endpoint)
        if view_func is not None:
            rate_limits: List[RateLimit] = getattr(view_func, QUART_RATE_LIMITER_ATTRIBUTE, [])
            await self._raise_on_rejection(endpoint, rate_limits)
            await self._update_limits(endpoint, rate_limits)

    async def _raise_on_rejection(self, endpoint: str, rate_limits: List[RateLimit]) -> None:
        now = datetime.utcnow()
        for rate_limit in rate_limits:
            key = await self._create_key(endpoint, rate_limit)
            # This is the GCRA rate limiting system and tat stands for
            # the theoretical arrival time.
            tat = max(await self.store.get(key, now), now)
            separation = (tat - now).total_seconds()
            max_interval = rate_limit.period.total_seconds() - rate_limit.inverse
            if separation > max_interval:
                retry_after = ((tat - timedelta(seconds=max_interval)) - now).total_seconds()
                raise RateLimitExceeded(int(retry_after))

    async def _update_limits(self, endpoint: str, rate_limits: List[RateLimit]) -> None:
        # Update the tats for all the rate limits. This must only
        # occur if no limit rejects the request.
        now = datetime.utcnow()
        for rate_limit in rate_limits:
            key = await self._create_key(endpoint, rate_limit)
            tat = max(await self.store.get(key, now), now)
            new_tat = max(tat, now) + timedelta(seconds=rate_limit.inverse)
            await self.store.set(key, new_tat)

    async def _after_request(self, response: Response) -> Response:
        endpoint = request.endpoint
        view_func = current_app.view_functions.get(endpoint)
        rate_limits: List[RateLimit] = getattr(view_func, QUART_RATE_LIMITER_ATTRIBUTE, [])
        try:
            min_limit = min(rate_limits, key=lambda rate_limit: rate_limit.period.total_seconds())
        except ValueError:
            pass  # No rate limits
        else:
            key = await self._create_key(endpoint, min_limit)
            now = datetime.utcnow()
            tat = max(await self.store.get(key, now), now)
            separation = (tat - now).total_seconds()
            remaining = int((min_limit.period.total_seconds() - separation) / min_limit.inverse)
            response.headers["RateLimit-Limit"] = str(min_limit.count)
            response.headers["RateLimit-Remaining"] = str(remaining)
            response.headers["RateLimit-Reset"] = str(int(separation))
        return response

    async def _create_key(self, endpoint: str, rate_limit: RateLimit) -> str:
        key_function = rate_limit.key_function or self.key_function
        key = await key_function()
        app_name = current_app.import_name
        return f"{app_name}-{endpoint}-{rate_limit.key}-{key}"
