"""
Rate limiting utilities for controlling execution frequency of functions and code blocks.

This module provides decorators, context managers, and utility functions to limit how
often code executes, which is particularly useful for:

- Reducing noisy output in loops (print statements, logging)
- Throttling expensive operations
- Controlling debug message frequency
- Implementing backpressure in data processing pipelines

The rate limiting tracks execution counts per unique key and shows suppressed message
counts similar to system loggers (e.g., "... 5 more similar messages suppressed").

Basic Usage:
    # Decorator - limit function execution
    @rate_limit(every_n=10)
    def debug_function():
        print("Debug info")

    # Context manager - limit code block execution
    for i in range(100):
        with rate_limited(every_n=10, key="progress") as should_execute:
            if should_execute:
                print(f"Progress: {i}")

    # Convenience functions
    for i in range(100):
        rate_limited_print(f"Processing item {i}", every_n=10)

Advanced Usage:
    # Custom rate limiter instances for isolation
    debug_limiter = RateLimiter()

    @rate_limit(every_n=5, limiter=debug_limiter)
    def isolated_debug():
        print("This uses a separate counter")

    # Rate-limited logging with suppression tracking
    import logging
    logger = logging.getLogger(__name__)

    for i in range(50):
        rate_limited_log(logger, logging.WARNING,
                        "High CPU usage", every_n=10)

Classes:
    RateLimiter: Core rate limiting logic with per-key counters

Functions:
    rate_limit: Decorator for rate-limiting function calls
    rate_limited: Context manager for rate-limiting code blocks
    rate_limited_print: Convenience function for rate-limited printing
    rate_limited_log: Convenience function for rate-limited logging

Author: Generated by Claude
License: MIT
"""

import time
import logging
from functools import wraps
from collections import defaultdict
from contextlib import contextmanager
from typing import Dict, Any, Optional, Callable


class RateLimiter:
    """A rate limiter that tracks execution counts and provides rate limiting functionality."""

    def __init__(self):
        self.counters: Dict[str, int] = defaultdict(int)
        self.last_executed: Dict[str, int] = defaultdict(int)
        self.suppressed_count: Dict[str, int] = defaultdict(int)

    def should_execute(self, key: str, every_n: int) -> tuple[bool, int]:
        """
        Check if an operation should be executed based on rate limiting.

        Returns:
            tuple: (should_execute: bool, suppressed_count: int)
        """
        self.counters[key] += 1

        # Execute on 1st call, then every every_n calls after that
        if self.counters[key] == 1 or (self.counters[key] - 1) % every_n == 0:
            # For the first call, no suppressed count to return
            if self.counters[key] == 1:
                return True, 0
            else:
                # For subsequent executions, return the suppressed count and reset it
                suppressed = self.suppressed_count[key]
                self.suppressed_count[key] = 0
                return True, suppressed
        else:
            self.suppressed_count[key] += 1
            return False, 0

    def reset(self, key: Optional[str] = None):
        """Reset counters for a specific key or all keys."""
        if key:
            self.counters[key] = 0
            self.last_executed[key] = 0
            self.suppressed_count[key] = 0
        else:
            self.counters.clear()
            self.last_executed.clear()
            self.suppressed_count.clear()


# Global rate limiter instance
_global_limiter = RateLimiter()


def rate_limit(
    every_n: int = 10, key: Optional[str] = None, show_suppressed: bool = True, limiter: Optional[RateLimiter] = None
):
    """
    Decorator that limits function execution to every N calls.

    Args:
        every_n: Execute only every N calls
        key: Unique identifier for this rate limiter (auto-generated if None)
        show_suppressed: Whether to show "... N more suppressed" message
        limiter: Custom RateLimiter instance (uses global one if None)
    """

    def decorator(func: Callable) -> Callable:
        nonlocal key
        if key is None:
            key = f"{func.__module__}.{func.__qualname__}"

        rate_limiter = limiter or _global_limiter

        @wraps(func)
        def wrapper(*args, **kwargs):
            should_execute, suppressed = rate_limiter.should_execute(key, every_n)

            if should_execute:
                # Show suppressed count BEFORE executing (except for first execution)
                if show_suppressed and suppressed > 0:
                    print(f"... {suppressed} more similar messages suppressed")

                result = func(*args, **kwargs)
                return result

            return None  # Suppressed execution

        # Add utility methods to the wrapper
        wrapper.reset = lambda: rate_limiter.reset(key)
        wrapper.get_count = lambda: rate_limiter.counters[key]

        return wrapper

    return decorator


@contextmanager
def rate_limited(
    every_n: int = 10, key: Optional[str] = None, show_suppressed: bool = True, limiter: Optional[RateLimiter] = None
):
    """
    Context manager that limits execution of code block to every N calls.

    Args:
        every_n: Execute only every N calls
        key: Unique identifier for this rate limiter (required for context manager)
        show_suppressed: Whether to show "... N more suppressed" message
        limiter: Custom RateLimiter instance (uses global one if None)

    Yields:
        bool: True if code should execute, False if suppressed
    """
    if key is None:
        raise ValueError("Context manager requires a unique 'key' parameter")

    rate_limiter = limiter or _global_limiter
    should_execute, suppressed = rate_limiter.should_execute(key, every_n)

    if should_execute:
        # Show suppressed count BEFORE yielding (except for first execution)
        if show_suppressed and suppressed > 0:
            print(f"... {suppressed} more similar operations suppressed")

    try:
        yield should_execute
    finally:
        pass  # Nothing to do in finally


# Convenience functions for common use cases
def rate_limited_print(message: str, every_n: int = 10, key: Optional[str] = None):
    """Print a message only every N calls."""
    if key is None:
        key = f"print_{hash(message) % 10000}"

    with rate_limited(every_n, key=key) as should_execute:
        if should_execute:
            print(message)


def rate_limited_log(
    logger: logging.Logger, level: int, message: str, every_n: int = 10, key: Optional[str] = None, *args
):
    """Log a message only every N calls."""
    if key is None:
        key = f"log_{hash(message) % 10000}"

    with rate_limited(every_n, key=key, show_suppressed=False) as should_execute:
        if should_execute:
            suppressed = _global_limiter.suppressed_count[key]
            if suppressed > 0:
                message += f" (... {suppressed} more similar messages suppressed)"
            logger.log(level, message, *args)


# Example usage and demonstrations
if __name__ == "__main__":
    print("=== Decorator Example ===")

    @rate_limit(every_n=3, show_suppressed=True)
    def noisy_function(i):
        print(f"Processing item {i}")
        return f"result_{i}"

    # This will only print every 3rd call
    for i in range(10):
        noisy_function(i)

    print(f"\nTotal calls to noisy_function: {noisy_function.get_count()}")

    print("\n=== Context Manager Example ===")

    # Using context manager in a loop
    for i in range(15):
        with rate_limited(every_n=5, key="loop_progress") as should_execute:
            if should_execute:
                print(f"Progress update: processed {i + 1} items")

    print("\n=== Convenience Functions Example ===")

    # Using convenience functions
    for i in range(8):
        rate_limited_print(f"Item {i} processed", every_n=3, key="item_processing")

    print("\n=== Logging Example ===")

    # Setup logging
    logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
    logger = logging.getLogger(__name__)

    # Rate-limited logging
    for i in range(12):
        rate_limited_log(
            logger, logging.WARNING, f"High memory usage detected: {i * 10}%", every_n=4, key="memory_warning"
        )

    print("\n=== Multiple Rate Limiters Example ===")

    # Using separate rate limiter instances for different purposes
    debug_limiter = RateLimiter()
    error_limiter = RateLimiter()

    @rate_limit(every_n=2, key="debug_messages", limiter=debug_limiter)
    def debug_message(msg):
        print(f"DEBUG: {msg}")

    @rate_limit(every_n=4, key="error_messages", limiter=error_limiter)
    def error_message(msg):
        print(f"ERROR: {msg}")

    for i in range(10):
        debug_message(f"Debug info {i}")
        if i % 2 == 0:
            error_message(f"Error condition {i}")
