"""Tiered per-route rate limiting with a shadow|enforce mode.

flask-limiter's automatic enforcement raises inside ``before_request`` and cannot let a
request continue, so it cannot implement a shadow (count-and-log, never block) mode. We
therefore drive enforcement ourselves: parse the applicable limit strings with the ``limits``
library, hit each against one shared sliding-window store keyed by ``principal_key()``, and on
breach either log only (shadow) or raise ``TooManyRequests`` (enforce) so the registered
handler renders a leak-free 429.

A view opts into a tighter or looser tier with ``@rate_category(...)``; untagged views fall back
to the per-key environmental default. A category may carry several limits, all enforced together,
so a sustained cap and a short burst cap coexist.

Every rate-limited response carries ``X-RateLimit-Limit/Remaining/Reset`` for the most
constraining window (fewest remaining, then soonest reset) so clients can self-throttle before
they breach rather than discovering the cap by getting a 429; the enforced 429 additionally
carries ``Retry-After``. Headers are published in shadow mode too, so clients can begin adapting
before enforcement is turned on.
"""

import logging
import time
from collections.abc import Callable
from enum import Enum
from typing import TypeVar

from flask import Flask, Response, g, request
from limits import RateLimitItem, parse, parse_many
from werkzeug.exceptions import TooManyRequests

from core.config import Config
from core.hardening import observability
from core.hardening.client_ip import principal_key
from core.hardening.rate_limit_storage import RateLimiterFactory, build_rate_limiter
from core.hardening.rate_policy import (
    RATE_CATEGORIES,
    RATE_EXEMPT,
    RateCategory,
)
from core.hardening.view_markers import resolve_view_marker

logger = logging.getLogger(__name__)

_CATEGORY_ATTR = '_rate_category'
_STATE_KEY = 'hardening_rate_limiter'

F = TypeVar('F', bound=Callable[..., object])


class RateLimitMode(str, Enum):
    """How the limiter reacts to a breach.

    There is deliberately no ``off`` value: protective middleware is never disableable, so a
    misconfigured mode must fail fast at startup rather than silently degrade to never-block.
    """

    SHADOW = 'shadow'  # count and log the breach, never block
    ENFORCE = 'enforce'  # raise 429 on breach


class _RateLimitState:
    """Per-app rate-limit state: the shared limiter factory plus the resolved limit buckets.

    Stored on ``app.extensions`` so it is isolated per app (each test app builds a fresh
    in-memory store) and reachable if a caller needs to clear it.
    """

    def __init__(
        self,
        factory: RateLimiterFactory,
        default_limits: list[RateLimitItem],
        category_limits: dict[RateCategory, list[RateLimitItem]],
        mode: RateLimitMode,
    ) -> None:
        """Hold the shared limiter factory, the default + per-category parsed limits, and the mode."""
        self.factory = factory
        self.default_limits = default_limits
        self.category_limits = category_limits
        self.mode = mode


def rate_category(category: RateCategory) -> Callable[[F], F]:
    """Tag a view with its rate-limit category. Returns the view unchanged (no wrapping)."""

    def decorator(view: F) -> F:
        setattr(view, _CATEGORY_ATTR, category)
        return view

    return decorator


def setup_rate_limiting(app: Flask, config: type[Config]) -> None:
    """Install a self-driven, per-route rate limiter on ``app``.

    Resolves each request's applicable limits (category limits if the view is
    ``@rate_category``-tagged, else the per-key default), hits every one against the shared
    sliding-window store keyed by the principal, publishes the most-constraining window as
    ``X-RateLimit-*`` headers, and on breach logs and -- only in enforce mode -- raises so the
    generic 429 handler responds (with ``Retry-After``).
    """
    try:
        mode = RateLimitMode(config.RATELIMIT_MODE)
    except ValueError:
        raise ValueError(
            f'RATELIMIT_MODE must be one of {[m.value for m in RateLimitMode]}, '
            f'got {config.RATELIMIT_MODE!r}'
        ) from None
    factory = build_rate_limiter(config.RATELIMIT_STORAGE_URI)
    default_limits = parse_many(
        config.RATELIMIT_DEFAULT
    )  # raises on empty/garbage (fail fast)
    category_limits = {
        category: [parse(spec) for spec in specs]
        for category, specs in RATE_CATEGORIES.items()
    }
    state = _RateLimitState(factory, default_limits, category_limits, mode)
    app.extensions[_STATE_KEY] = state

    @app.before_request
    def _enforce_rate_limit() -> None:
        endpoint = request.endpoint
        if not endpoint or endpoint == 'static' or endpoint in RATE_EXEMPT:
            return

        # resolve_view_marker resolves @rate_category on class-based MethodView routes,
        # where the marker lands on the get()/post() method rather than the as_view()
        # closure in view_functions (with Flask's HEAD->get fallback).
        category: RateCategory | None = resolve_view_marker(
            app, endpoint, request.method, _CATEGORY_ATTR, None
        )
        limits = state.category_limits[category] if category else state.default_limits

        key, key_type = principal_key()
        # Hit every applicable limit (do not short-circuit): in shadow mode each window must
        # advance independently so the counts faithfully predict what enforce mode would do
        # and the published headers stay accurate. Indexed so the breached window's own stats
        # drive Retry-After.
        limiter = state.factory.for_request()
        admitted = [limiter.hit(item, key) for item in limits]
        stats = [limiter.get_window_stats(item, key) for item in limits]

        # Publish the most-constraining window (fewest remaining, then soonest reset) so a
        # client can throttle before it breaches. WindowStats is (reset_time, remaining).
        governing = min(
            range(len(limits)),
            key=lambda i: (stats[i].remaining, stats[i].reset_time),
        )
        g.rate_limit_limit = limits[governing].amount
        g.rate_limit_remaining = stats[governing].remaining
        g.rate_limit_reset = int(stats[governing].reset_time)

        breached = [i for i, ok in enumerate(admitted) if not ok]
        if not breached:
            return

        # Log and count the breach with the key TYPE only (the key value is PII and
        # high-cardinality). The counter fires in shadow too, so the rate of would-be-429s per
        # category/key_type is observable in Datadog before enforcement is turned on.
        category_name = category.value if category else 'default'
        logger.warning(
            'rate limit breach endpoint=%s category=%s key_type=%s mode=%s',
            endpoint,
            category_name,
            key_type,
            state.mode.value,
        )
        observability.rate_limit_rejected(category_name, key_type)
        if state.mode is RateLimitMode.ENFORCE:
            # Retry-After is the LONGEST reset among ALL breached windows, so a client is never
            # told to retry before the slowest cap that blocked it has cleared (a burst can trip
            # both the per-second and per-minute windows at once).
            now = time.time()
            g.rate_limit_retry_after = max(
                1, max(int(stats[i].reset_time - now) + 1 for i in breached)
            )
            raise TooManyRequests

    @app.after_request
    def _set_rate_limit_headers(response: Response) -> Response:
        limit = g.get('rate_limit_limit')
        if limit is None:
            return response
        response.headers['X-RateLimit-Limit'] = str(limit)
        response.headers['X-RateLimit-Remaining'] = str(
            g.get('rate_limit_remaining', 0)
        )
        response.headers['X-RateLimit-Reset'] = str(g.get('rate_limit_reset', 0))
        retry_after = g.get('rate_limit_retry_after')
        if retry_after is not None:
            response.headers['Retry-After'] = str(retry_after)
        return response