"""Rate limiting module using token bucket algorithm.

This module provides rate limiting capabilities for job processing,
preventing any single job type from overwhelming the system.

Example:
    >>> from oclawma.ratelimit import RateLimiter, RateLimitConfig
    >>> config = RateLimitConfig(rate=10.0, burst=5)
    >>> limiter = RateLimiter(config)
    >>> info = limiter.try_acquire("email")
    >>> if info.allowed:
    ...     process_job()
    ... else:
    ...     retry_after = info.retry_after
"""

from __future__ import annotations

import threading
import time
from dataclasses import dataclass, field
from typing import Literal


@dataclass
class RateLimitInfo:
    """Information about rate limit status.

    Attributes:
        allowed: Whether the request was allowed
        tokens_remaining: Number of tokens remaining in bucket
        tokens_total: Total capacity of the bucket (burst)
        reset_after: Seconds until bucket is fully refilled
        retry_after: Seconds to wait before retry (0 if allowed)
        job_type: The job type this info applies to
    """

    allowed: bool = True
    tokens_remaining: int = 0
    tokens_total: int = 0
    reset_after: float = 0.0
    retry_after: float = 0.0
    job_type: str = "default"

    def to_headers(self) -> dict[str, str]:
        """Convert to HTTP-style headers dictionary.

        Returns:
            Dictionary with rate limit headers
        """
        return {
            "X-RateLimit-Limit": str(self.tokens_total),
            "X-RateLimit-Remaining": str(self.tokens_remaining),
            "X-RateLimit-Reset": f"{self.reset_after:.3f}",
            "X-RateLimit-Retry-After": f"{self.retry_after:.3f}",
            "X-RateLimit-JobType": self.job_type,
        }


@dataclass
class RateLimitConfig:
    """Configuration for rate limiting.

    Attributes:
        rate: Tokens added per second (refill rate)
        burst: Maximum number of tokens in bucket
        per_job_type: Per-job-type rate limit overrides
        enabled: Whether rate limiting is enabled
        default_behavior: Default behavior when rate limited ("block", "drop", "defer")
    """

    rate: float = 10.0
    burst: int = 5
    per_job_type: dict[str, RateLimitConfig] = field(default_factory=dict)
    enabled: bool = True
    default_behavior: Literal["block", "drop", "defer"] = "block"

    def to_dict(self) -> dict:
        """Convert config to dictionary."""
        return {
            "rate": self.rate,
            "burst": self.burst,
            "enabled": self.enabled,
            "default_behavior": self.default_behavior,
            "per_job_type": {
                k: v.to_dict() if isinstance(v, RateLimitConfig) else v
                for k, v in self.per_job_type.items()
            },
        }

    @classmethod
    def from_dict(cls, data: dict) -> RateLimitConfig:
        """Create config from dictionary."""
        per_job_type = {}
        for key, value in data.get("per_job_type", {}).items():
            if isinstance(value, dict):
                per_job_type[key] = cls.from_dict(value)
            else:
                per_job_type[key] = value

        return cls(
            rate=data.get("rate", 10.0),
            burst=data.get("burst", 5),
            enabled=data.get("enabled", True),
            default_behavior=data.get("default_behavior", "block"),
            per_job_type=per_job_type,
        )


class RateLimitError(Exception):
    """Exception raised when rate limit is exceeded.

    Attributes:
        message: Error message
        info: RateLimitInfo with details about the limit
    """

    def __init__(self, message: str, info: RateLimitInfo):
        super().__init__(message)
        self.info = info


class TokenBucket:
    """Token bucket for rate limiting.

    Thread-safe implementation of the token bucket algorithm.
    Tokens are added at a constant rate up to a maximum burst.

    Attributes:
        rate: Tokens added per second
        burst: Maximum number of tokens
        tokens: Current number of tokens
        total_requests: Total number of requests made
        blocked_requests: Number of requests that were blocked
    """

    def __init__(self, rate: float, burst: int, initial_tokens: int | None = None) -> None:
        """Initialize the token bucket.

        Args:
            rate: Tokens added per second
            burst: Maximum number of tokens
            initial_tokens: Starting number of tokens (default: burst)
        """
        self.rate = rate
        self.burst = burst
        self.tokens = initial_tokens if initial_tokens is not None else burst
        self._last_refill = time.monotonic()
        self._lock = threading.Lock()
        self.total_requests = 0
        self.blocked_requests = 0

    def _refill(self) -> None:
        """Refill tokens based on elapsed time."""
        now = time.monotonic()
        elapsed = now - self._last_refill
        tokens_to_add = elapsed * self.rate

        self.tokens = min(self.burst, self.tokens + tokens_to_add)
        self._last_refill = now

    def consume(self, tokens: int = 1) -> bool:
        """Try to consume tokens from the bucket.

        Args:
            tokens: Number of tokens to consume

        Returns:
            True if tokens were consumed, False otherwise
        """
        with self._lock:
            self._refill()
            self.total_requests += 1

            if tokens <= 0:
                return True

            if self.tokens >= tokens:
                self.tokens -= tokens
                return True

            self.blocked_requests += 1
            return False

    def get_stats(self) -> dict:
        """Get current bucket statistics.

        Returns:
            Dictionary with bucket statistics
        """
        with self._lock:
            self._refill()
            return {
                "tokens": int(self.tokens),
                "burst": self.burst,
                "rate": self.rate,
                "total_requests": self.total_requests,
                "blocked_requests": self.blocked_requests,
                "utilization": 1.0 - (self.tokens / self.burst) if self.burst > 0 else 0.0,
            }

    def reset(self) -> None:
        """Reset the bucket to full."""
        with self._lock:
            self.tokens = self.burst
            self._last_refill = time.monotonic()
            self.total_requests = 0
            self.blocked_requests = 0

    def time_to_refill(self, tokens: int = 1) -> float:
        """Calculate time until enough tokens are available.

        Args:
            tokens: Number of tokens needed

        Returns:
            Seconds until tokens are available (0 if available now)
        """
        with self._lock:
            self._refill()

            if self.tokens >= tokens:
                return 0.0

            tokens_needed = tokens - self.tokens
            return tokens_needed / self.rate if self.rate > 0 else float("inf")


class RateLimiter:
    """Rate limiter for job processing.

    Manages token buckets for different job types and provides
    methods to acquire tokens with various behaviors.

    Attributes:
        config: RateLimitConfig with rate limiting settings
        _buckets: Dictionary of job type to TokenBucket
        _lock: Threading lock for bucket creation
    """

    def __init__(self, config: RateLimitConfig | None = None) -> None:
        """Initialize the rate limiter.

        Args:
            config: RateLimitConfig instance (default: default config)
        """
        self.config = config or RateLimitConfig()
        self._buckets: dict[str, TokenBucket] = {}
        self._lock = threading.Lock()

        # Create default bucket
        if self.config.enabled:
            self._buckets["default"] = TokenBucket(
                rate=self.config.rate,
                burst=self.config.burst,
            )

    def _get_bucket(self, job_type: str) -> TokenBucket:
        """Get or create a token bucket for a job type.

        Args:
            job_type: The job type to get bucket for

        Returns:
            TokenBucket for the job type
        """
        # Fast path: bucket exists
        if job_type in self._buckets:
            return self._buckets[job_type]

        # Slow path: create bucket (needs lock)
        with self._lock:
            # Double-check after acquiring lock
            if job_type in self._buckets:
                return self._buckets[job_type]

            # Check for per-job-type config
            if job_type in self.config.per_job_type:
                job_config = self.config.per_job_type[job_type]
                bucket = TokenBucket(
                    rate=job_config.rate,
                    burst=job_config.burst,
                )
            else:
                # Use default config
                bucket = TokenBucket(
                    rate=self.config.rate,
                    burst=self.config.burst,
                )

            self._buckets[job_type] = bucket
            return bucket

    def try_acquire(
        self,
        job_type: str,
        tokens: int = 1,
    ) -> RateLimitInfo:
        """Try to acquire tokens without blocking.

        Args:
            job_type: The job type to acquire tokens for
            tokens: Number of tokens to acquire (default: 1)

        Returns:
            RateLimitInfo with result and quota information
        """
        if not self.config.enabled:
            return RateLimitInfo(
                allowed=True,
                tokens_remaining=-1,
                tokens_total=-1,
                job_type=job_type,
            )

        bucket = self._get_bucket(job_type)
        allowed = bucket.consume(tokens)

        stats = bucket.get_stats()
        retry_after = bucket.time_to_refill(tokens) if not allowed else 0.0
        reset_after = (bucket.burst - stats["tokens"]) / bucket.rate if bucket.rate > 0 else 0.0

        return RateLimitInfo(
            allowed=allowed,
            tokens_remaining=int(stats["tokens"]),
            tokens_total=bucket.burst,
            reset_after=reset_after,
            retry_after=retry_after,
            job_type=job_type,
        )

    def acquire(
        self,
        job_type: str,
        tokens: int = 1,
        blocking: bool = True,
        timeout: float | None = None,
    ) -> RateLimitInfo:
        """Acquire tokens, optionally blocking until available.

        Args:
            job_type: The job type to acquire tokens for
            tokens: Number of tokens to acquire (default: 1)
            blocking: Whether to block until tokens available (default: True)
            timeout: Maximum seconds to block (None = forever)

        Returns:
            RateLimitInfo with result and quota information
        """
        start_time = time.monotonic()

        while True:
            info = self.try_acquire(job_type, tokens)

            if info.allowed or not blocking:
                return info

            # Check timeout
            if timeout is not None:
                elapsed = time.monotonic() - start_time
                if elapsed >= timeout:
                    return info

            # Wait a bit before retrying (exponential backoff up to 100ms)
            wait_time = min(info.retry_after / 10, 0.1)
            if wait_time <= 0:
                wait_time = 0.01

            # Don't wait longer than remaining timeout
            if timeout is not None:
                remaining = timeout - (time.monotonic() - start_time)
                wait_time = min(wait_time, remaining)
                if wait_time <= 0:
                    return info

            time.sleep(wait_time)

    def get_quota(self, job_type: str) -> RateLimitInfo:
        """Get current quota without consuming tokens.

        Args:
            job_type: The job type to check quota for

        Returns:
            RateLimitInfo with current quota (always allowed=True)
        """
        if not self.config.enabled:
            return RateLimitInfo(
                allowed=True,
                tokens_remaining=-1,
                tokens_total=-1,
                job_type=job_type,
            )

        bucket = self._get_bucket(job_type)
        stats = bucket.get_stats()

        return RateLimitInfo(
            allowed=True,
            tokens_remaining=stats["tokens"],
            tokens_total=bucket.burst,
            reset_after=(bucket.burst - stats["tokens"]) / bucket.rate if bucket.rate > 0 else 0.0,
            retry_after=0.0,
            job_type=job_type,
        )

    def reset(self, job_type: str | None = None) -> None:
        """Reset rate limiter state.

        Args:
            job_type: Specific job type to reset, or None for all
        """
        with self._lock:
            if job_type is not None:
                if job_type in self._buckets:
                    self._buckets[job_type].reset()
            else:
                for bucket in self._buckets.values():
                    bucket.reset()

    def get_stats(self) -> dict[str, dict]:
        """Get statistics for all buckets.

        Returns:
            Dictionary mapping job type to bucket statistics
        """
        return {job_type: bucket.get_stats() for job_type, bucket in self._buckets.items()}
