"""
federation/rate_limit.py
Federation rate limiting implementation.

Features:
- Per-domain rate limiting
- Multiple rate limit strategies
- In-memory storage
- Configurable limits
- Burst handling
"""

from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from enum import Enum
from http import HTTPStatus
from logging import Logger
from urllib.parse import urlparse

from fastapi.responses import JSONResponse

from phederation.cache import BaseCache
from phederation.utils.base import ObjectId, validate_object_id, RateLimit
from phederation.utils.logging import configure_logger
from phederation.utils.settings import PhedSettings
from phederation.utils.version import PHEDERATION_HEADERS


class RateLimitStrategy(Enum):
    """Rate limit strategies."""

    FIXED_WINDOW = "fixed_window"
    SLIDING_WINDOW = "sliding_window"
    TOKEN_BUCKET = "token_bucket"
    LEAKY_BUCKET = "leaky_bucket"


@dataclass
class RateLimitState:
    """Current rate limit state for a domain.

    Counts how many requests are `remaining` until a certain `reset` date.
    """

    remaining: int
    """ Number of remaining requests during a set period. """

    reset: datetime
    """ The time to reset the remaining requests counter. """

    limit: int
    """ The maximum for this rate limit state, to reset to at `reset` time. Set from the general `RateLimit`. """


class RateLimiter:
    """Federation rate limiting."""

    EVERYBODY: ObjectId = validate_object_id(url="http://everybody.localhost")
    """ The generic id to identify a non-identified domain. This rate limits "everybody", i.e. sets the limit on requests on non-identified urls. """

    def __init__(
        self,
        cache: BaseCache,
        settings: PhedSettings,
        strip_domains: bool = True,
    ):
        """Initialize rate limiter.

        Args:
            cache: The cache to use for storing urls
            default_limit: Default rate limit settings
            strategy: Rate limiting strategy
            strip_domains: If input domains should be stripped before caching. This can be used to either rate limit entire domains or specific actors.
        """
        self.settings: PhedSettings = settings
        # TODO: implement rate limiting strategy and load from settings
        self.strategy: RateLimitStrategy = RateLimitStrategy.FIXED_WINDOW
        self.policies: dict[str, RateLimit] = self.settings.rate_limiter.policies
        self.default_policy: RateLimit = self.settings.rate_limiter.default_policy
        self.cache: BaseCache = cache
        self.strip_domains: bool = True
        self.logger: Logger = configure_logger(__name__, prefix=self.settings.federation.logging_prefix)

    def _strip_domain(self, domain: str):
        domain = domain if "://" in domain else f"https://{domain}"
        if self.strip_domains:
            parsed_domain = urlparse(domain)
            return parsed_domain.netloc
        else:
            return domain

    def _get_policy(self, domain: str) -> RateLimit:
        if policy := self.policies.get(domain, None):
            return policy
        return self.default_policy

    def _get_state(self, domain: str) -> RateLimitState:
        """Get current rate limit state for domain."""
        domain = self._strip_domain(domain)
        state = self.cache.get(f"ratelimit:{domain}")
        if not state:
            policy = self._get_policy(domain)
            state = RateLimitState(
                remaining=policy.requests,
                reset=datetime.now(timezone.utc) + timedelta(seconds=policy.period),
                limit=policy.requests,
            )
            self._set_state(domain, state)
        return state

    def _set_state(self, domain: str, state: RateLimitState) -> None:
        """Set rate limit state for domain."""
        domain = self._strip_domain(domain)
        self.cache.set(f"ratelimit:{domain}", state)

    def ratelimit_response(self, domain: str) -> JSONResponse:
        headers = self.update_headers(domain=domain, headers=PHEDERATION_HEADERS)
        return JSONResponse(headers=headers, content={"status": "error", "detail": "rate limit reached"}, status_code=HTTPStatus.TOO_MANY_REQUESTS)

    def check_rate_limit(self, domain: str, reduce_counter: bool = True, strip_domain: bool = True) -> bool:
        """Check limit for given domain, and reduce request counter by one."""
        self.strip_domains = strip_domain
        domain = self._strip_domain(domain)
        state = self._get_state(domain)

        # Reset if expired
        if datetime.now(timezone.utc) >= state.reset:
            policy = self._get_policy(domain)
            state = RateLimitState(
                remaining=policy.requests,
                reset=datetime.now(timezone.utc) + timedelta(seconds=policy.period),
                limit=policy.requests,
            )
            self._set_state(domain, state)
            self.strip_domains = True
            return True

        if state.remaining <= 0:
            self.strip_domains = True
            return False

        # Update remaining
        if reduce_counter:
            state.remaining -= 1
            self._set_state(domain, state)
        self.strip_domains = True
        return True

    def get_wait_time(self, domain: str) -> float:
        """Get wait time in seconds until rate limit reset."""
        domain = self._strip_domain(domain)
        state = self._get_state(domain)
        if state.remaining > 0:
            return 0

        wait_time = (state.reset - datetime.now(timezone.utc)).total_seconds()
        return max(0, wait_time)

    def update_rate_limit(self, domain: str, headers: dict[str, str]) -> RateLimitState | None:
        """Update rate limit state from response headers.

        This allows remote instances to update the rate limits of themselves.
        Will be limited below by one to always be able to send at least one request.

        Header keys/values
         - X-RateLimit-Remaining (int): number of remaining requests that the remote instance accepts.
         - X-RateLimit-Reset (int): the reset timestamp when the remote request counter resets.
         - X-RateLimit-Limit (int): the total limit of requests accepted by the remote instance until the reset date.
        """
        domain = self._strip_domain(domain)
        headers_lower = {key.lower(): headers[key] for key in headers.keys()}
        remaining = headers_lower.get("x-ratelimit-remaining")
        reset = headers_lower.get("x-ratelimit-reset")
        limit = headers_lower.get("x-ratelimit-limit")

        if remaining is not None and reset is not None and limit is not None:
            state = RateLimitState(remaining=int(remaining), reset=datetime.fromtimestamp(float(reset), tz=timezone.utc), limit=max(1, int(limit)))
            self._set_state(domain, state)
            return state

    def update_headers(self, domain: str, headers: dict[str, str]):
        domain = self._strip_domain(domain)
        state = self._get_state(domain=domain)
        headers["x-ratelimit-remaining"] = str(state.remaining)
        headers["x-ratelimit-reset"] = str(state.reset.timestamp())
        headers["x-ratelimit-limit"] = str(state.limit)
        if state.remaining == 0:
            headers["retry-after"] = str(int((state.reset - datetime.now(timezone.utc)).total_seconds() + 1))
        return headers

    def clear(self, domain: str) -> None:
        domain = self._strip_domain(domain)
        """Clear rate limit state for domain."""
        self.cache.delete(f"ratelimit:{domain}")
