"""
Parallelism and Rate Limiting Manager for Task Processing

This module provides resource-level concurrency control and rate limiting for processing
entities/items in a controlled manner.

Key Concepts:
- parallelism_max: Maximum number of entities being processed simultaneously at any moment
- parallelism_throttle: Maximum rate (entities/second) at which new processing can start

Classes:
- ThreadExecutor: For I/O-bound tasks (API calls, file I/O) - uses ThreadPoolExecutor
- LoopExecutor: For CPU/GPU-bound tasks (OCR, models) - uses sequential for-loop
- BatchExecutor: For batch operations (NER, inference) - processes in batches

Example:
    >>> # I/O-bound task
    >>> config = ParallelismConfig(max_parallel=4, throttle=10)
    >>> executor = ThreadExecutor(config)
    >>> results = executor.process_items(files, extract_text)
    >>> 
    >>> # CPU/GPU-bound task
    >>> processor = LoopExecutor(config)
    >>> results = processor.process_items(images, ocr_process)
"""

import time
import threading
import logging
from dataclasses import dataclass
from typing import Optional, List, Callable, Any, TypeVar
from concurrent.futures import ThreadPoolExecutor, as_completed

logger = logging.getLogger(__name__)

T = TypeVar('T')
R = TypeVar('R')


@dataclass
class ParallelismConfig:
    """
    Configuration for controlling parallelism and throughput.
    
    Attributes:
        max_parallel: Maximum number of entities being processed simultaneously.
                     Acts as a concurrency limit (semaphore).
                     None = unlimited concurrency
        
        throttle: Maximum rate of entities starting processing per second.
                 Acts as a rate limiter.
                 None = unlimited rate
    
    Examples:
        >>> # At most 4 entities processed at once, no rate limit
        >>> config = ParallelismConfig(max_parallel=4, throttle=None)
        
        >>> # Unlimited concurrency, but max 10 entities/second
        >>> config = ParallelismConfig(max_parallel=None, throttle=10)
        
        >>> # Max 2 concurrent, max 5/second rate
        >>> config = ParallelismConfig(max_parallel=2, throttle=5)
    """
    max_parallel: Optional[int] = None
    throttle: Optional[int] = None
    
    def __post_init__(self):
        """Validate configuration values"""
        if self.max_parallel is not None and self.max_parallel < 1:
            raise ValueError(f"max_parallel must be >= 1, got {self.max_parallel}")
        if self.throttle is not None and self.throttle < 1:
            raise ValueError(f"throttle must be >= 1, got {self.throttle}")
    
    def has_limits(self) -> bool:
        """Check if any limits are configured"""
        return self.max_parallel is not None or self.throttle is not None
    
    def should_use_threading(self) -> bool:
        """
        Determine if threading should be used based on config.
        
        Threading only makes sense if max_parallel > 1.
        """
        return self.max_parallel is not None and self.max_parallel > 1


class RateLimiter:
    """
    Token bucket rate limiter.
    
    Allows up to `rate` operations per second, smoothing out bursts.
    
    Attributes:
        rate: Maximum operations per second
    
    Example:
        >>> limiter = RateLimiter(rate=10)
        >>> limiter.acquire()  # Blocks until a token is available
    """
    
    def __init__(self, rate: int):
        """
        Initialize rate limiter.
        
        Args:
            rate: Maximum operations per second (must be >= 1)
        
        Raises:
            ValueError: If rate < 1
        """
        if rate < 1:
            raise ValueError(f"rate must be >= 1, got {rate}")
        
        self.rate = rate
        self.tokens = 1.0  # Start with 1 token, not full bucket
        self.last_update = time.time()
        self.lock = threading.Lock()
    
    def acquire(self) -> None:
        """
        Acquire a token, blocking until one is available.
        
        Uses token bucket algorithm: tokens regenerate at `rate` per second,
        up to a maximum of `rate` tokens.
        """
        while True:
            with self.lock:
                now = time.time()
                elapsed = now - self.last_update
                
                # Add tokens based on elapsed time
                self.tokens = min(self.rate, self.tokens + elapsed * self.rate)
                self.last_update = now
                
                if self.tokens >= 1.0:
                    self.tokens -= 1.0
                    return
            
            # Wait a bit before trying again
            # Sleep time is calculated to minimize busy waiting
            time.sleep(1.0 / (self.rate * 10))


class ParallelismController:
    """
    Controls concurrency and rate limiting for entity processing.
    
    This ensures:
    1. No more than `max_parallel` entities are processing at any moment (if set)
    2. Entities start processing at most `throttle` per second (if set)
    
    Works with any execution model: for-loops, ThreadPools, async, etc.
    
    Example:
        >>> config = ParallelismConfig(max_parallel=4, throttle=10)
        >>> controller = ParallelismController(config)
        >>> 
        >>> # Use as context manager
        >>> with controller:
        ...     process_entity()  # Automatically acquires/releases
        >>> 
        >>> # Or manually
        >>> controller.acquire()
        >>> try:
        ...     process_entity()
        >>> finally:
        ...     controller.release()
    """
    
    def __init__(self, config: ParallelismConfig):
        """
        Initialize controller with given config.
        
        Args:
            config: ParallelismConfig instance
        """
        self.config = config
        self._semaphore: Optional[threading.Semaphore] = None
        self._rate_limiter: Optional[RateLimiter] = None
        
        # Initialize semaphore for concurrency control
        if config.max_parallel is not None:
            self._semaphore = threading.Semaphore(config.max_parallel)
            logger.debug(f"Initialized semaphore with {config.max_parallel} permits")
        
        # Initialize rate limiter
        if config.throttle is not None:
            self._rate_limiter = RateLimiter(config.throttle)
            logger.debug(f"Initialized rate limiter at {config.throttle}/second")
    
    def acquire(self) -> None:
        """
        Acquire permission to process an entity.
        
        Blocks until:
        1. A concurrency slot is available (if max_parallel is set)
        2. Rate limit allows (if throttle is set)
        
        The order matters: we acquire semaphore first, then wait for rate limit.
        This prevents holding a semaphore slot while waiting for rate limit.
        """
        # Wait for rate limit first (doesn't hold resources)
        if self._rate_limiter:
            self._rate_limiter.acquire()
        
        # Then acquire semaphore (holds a concurrency slot)
        if self._semaphore:
            self._semaphore.acquire()
    
    def release(self) -> None:
        """
        Release concurrency slot.
        
        Note: Only semaphore needs releasing; rate limiter is one-way.
        """
        if self._semaphore:
            self._semaphore.release()
    
    def __enter__(self):
        """Context manager entry - acquire permits"""
        self.acquire()
        return self
    
    def __exit__(self, *args):
        """Context manager exit - release permits"""
        self.release()


class ThreadExecutor:
    """
    Executes items using ThreadPoolExecutor with parallelism and rate limiting.
    
    Use for I/O-bound tasks (API calls, file I/O, network requests).
    Respects both max_parallel (concurrency limit) and throttle (rate limit).
    
    Example:
        >>> config = ParallelismConfig(max_parallel=4, throttle=10)
        >>> executor = ThreadExecutor(config)
        >>> 
        >>> def fetch_data(url):
        ...     return requests.get(url).text
        >>> 
        >>> results = executor.process_items(urls, fetch_data)
    """
    
    def __init__(self, config: ParallelismConfig):
        """
        Initialize thread executor with given config.
        
        Args:
            config: ParallelismConfig instance
        """
        self.config = config
        self.controller = ParallelismController(config)
    
    def process_items(
        self, 
        items: List[T], 
        process_fn: Callable[[T], R]
    ) -> List[Optional[R]]:
        """
        Process items using ThreadPoolExecutor with parallelism control.
        
        Args:
            items: List of items to process
            process_fn: Function to apply to each item. Should take one item
                       and return a result. Exceptions are caught and logged.
        
        Returns:
            List of results in same order as items. Failed items return None.
        
        Example:
            >>> # I/O-bound task (Tika text extraction)
            >>> results = executor.process_items(files, extract_text_with_tika)
        """
        if not items:
            return []
        
        # If no limits configured, use fastest method
        if not self.config.has_limits():
            logger.debug(f"Processing {len(items)} items with ThreadPool (no limits)")
            return self._process_unlimited(items, process_fn)
        
        # With limits
        logger.debug(
            f"Processing {len(items)} items with ThreadPool: "
            f"max_parallel={self.config.max_parallel}, "
            f"throttle={self.config.throttle}"
        )
        return self._process_limited(items, process_fn)
    
    def _process_unlimited(
        self, 
        items: List[T], 
        process_fn: Callable[[T], R]
    ) -> List[Optional[R]]:
        """
        Process with ThreadPool without limits.
        
        Uses a reasonable default thread count based on item count.
        """
        # Use reasonable thread count (not too many)
        max_workers = min(len(items), 10)
        
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            futures = [executor.submit(process_fn, item) for item in items]
            results = []
            for i, future in enumerate(futures):
                try:
                    results.append(future.result())
                except Exception as e:
                    logger.error(f"Processing error on item {i+1}: {e}", exc_info=True)
                    results.append(None)
        return results
    
    def _process_limited(
        self, 
        items: List[T], 
        process_fn: Callable[[T], R]
    ) -> List[Optional[R]]:
        """
        Process with ThreadPool respecting limits.
        
        Each thread acquires permits before processing, ensuring we respect
        both concurrency and rate limits across all threads.
        """
        # Use parallelism_max as thread pool size
        max_workers = self.config.max_parallel or 5
        
        def wrapped_fn(item_with_idx):
            idx, item = item_with_idx
            # Each thread acquires permits before processing
            with self.controller:
                try:
                    result = process_fn(item)
                    logger.debug(f"Thread processed item {idx+1}/{len(items)}")
                    return (idx, result)
                except Exception as e:
                    logger.error(f"Processing error on item {idx+1}: {e}", exc_info=True)
                    return (idx, None)
        
        # Initialize results list with None
        results = [None] * len(items)
        
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            # Submit all tasks with their indices
            futures = [
                executor.submit(wrapped_fn, (idx, item)) 
                for idx, item in enumerate(items)
            ]
            
            # Collect results as they complete
            for future in as_completed(futures):
                idx, result = future.result()
                results[idx] = result
        
        return results


class LoopExecutor:
    """
    Executes items sequentially using a for-loop with rate limiting.
    
    Use for CPU/GPU-bound tasks (OCR, model inference) where parallelism
    doesn't help or hurts performance (e.g., single GPU).
    
    Only respects throttle (rate limit), not max_parallel (already sequential).
    
    Example:
        >>> config = ParallelismConfig(max_parallel=None, throttle=10)
        >>> processor = LoopExecutor(config)
        >>> 
        >>> def ocr_image(image):
        ...     return easyocr_reader.readtext(image)
        >>> 
        >>> results = processor.process_items(images, ocr_image)
    """
    
    def __init__(self, config: ParallelismConfig):
        """
        Initialize loop processor with given config.
        
        Args:
            config: ParallelismConfig instance (only throttle is used)
        """
        self.config = config
        # Create controller with only rate limiting (no semaphore for sequential)
        rate_limit_only_config = ParallelismConfig(
            max_parallel=None,  # Ignore concurrency for sequential processing
            throttle=config.throttle
        )
        self.controller = ParallelismController(rate_limit_only_config)
    
    def process_items(
        self, 
        items: List[T], 
        process_fn: Callable[[T], R]
    ) -> List[Optional[R]]:
        """
        Process items sequentially with rate limiting.
        
        Args:
            items: List of items to process
            process_fn: Function to apply to each item. Should take one item
                       and return a result. Exceptions are caught and logged.
        
        Returns:
            List of results in same order as items. Failed items return None.
        
        Example:
            >>> # GPU-bound task (OCR)
            >>> results = processor.process_items(images, ocr_process)
        """
        if not items:
            return []
        
        # If no rate limit configured, use fastest method
        if self.config.throttle is None:
            logger.debug(f"Processing {len(items)} items sequentially (no rate limit)")
            return self._process_unlimited(items, process_fn)
        
        # With rate limiting
        logger.debug(
            f"Processing {len(items)} items sequentially with rate limit: "
            f"throttle={self.config.throttle}/sec"
        )
        return self._process_limited(items, process_fn)
    
    def _process_unlimited(
        self, 
        items: List[T], 
        process_fn: Callable[[T], R]
    ) -> List[Optional[R]]:
        """Process items sequentially without limits (fastest path)"""
        results = []
        for i, item in enumerate(items):
            try:
                results.append(process_fn(item))
            except Exception as e:
                logger.error(f"Processing error on item {i+1}: {e}", exc_info=True)
                results.append(None)
        return results
    
    def _process_limited(
        self, 
        items: List[T], 
        process_fn: Callable[[T], R]
    ) -> List[Optional[R]]:
        """
        Process items sequentially with rate limiting.
        
        Each item acquires permits before processing, ensuring we never
        exceed the rate limit.
        """
        results = []
        for i, item in enumerate(items):
            # Acquire rate limit permit (blocks if necessary)
            with self.controller:
                try:
                    result = process_fn(item)
                    results.append(result)
                    logger.debug(f"Processed item {i+1}/{len(items)}")
                except Exception as e:
                    logger.error(f"Processing error on item {i+1}: {e}", exc_info=True)
                    results.append(None)
        return results


class BatchExecutor:
    """
    Process items in batches with parallelism control.
    
    Useful for model inference where batching improves throughput
    (e.g., neural network inference, OCR, NER).
    
    Example:
        >>> config = ParallelismConfig(max_parallel=None, throttle=10)
        >>> processor = BatchExecutor(config)
        >>> 
        >>> def process_batch(texts):
        ...     # Process batch of texts together
        ...     return model.predict(texts)
        >>> 
        >>> results = processor.process_batches(
        ...     all_texts,
        ...     process_batch,
        ...     batch_size=10
        ... )
    """
    
    def __init__(self, config: ParallelismConfig):
        """
        Initialize batch processor with given config.
        
        Args:
            config: ParallelismConfig instance
        """
        self.config = config
        self.controller = ParallelismController(config)
    
    def process_batches(
        self, 
        items: List[T], 
        batch_fn: Callable[[List[T]], List[R]],
        batch_size: int = 1
    ) -> List[Optional[R]]:
        """
        Process items in batches.
        
        Args:
            items: List of items to process
            batch_fn: Function that takes a batch (list) and returns a list of results.
                     The function should return results in the same order as input.
                     Length of output must match length of input batch.
            batch_size: Size of each batch. Should be tuned based on model capacity
                       and memory constraints.
        
        Returns:
            List of results in same order as items. Failed items return None.
        
        Note:
            The entire batch is processed atomically (acquires permits once per batch,
            not once per item), since batching assumes the batch is processed together.
        
        Example:
            >>> # Process texts in batches of 10
            >>> def ner_batch(texts):
            ...     return ner_model.process_batch(texts)
            >>> 
            >>> results = processor.process_batches(
            ...     texts,
            ...     ner_batch,
            ...     batch_size=10
            ... )
        """
        if not items:
            return []
        
        if batch_size < 1:
            raise ValueError(f"batch_size must be >= 1, got {batch_size}")
        
        results = []
        num_batches = (len(items) + batch_size - 1) // batch_size
        
        logger.debug(
            f"Processing {len(items)} items in {num_batches} batches "
            f"(batch_size={batch_size})"
        )
        
        for batch_idx in range(0, len(items), batch_size):
            batch = items[batch_idx:batch_idx + batch_size]
            batch_num = (batch_idx // batch_size) + 1
            
            # Acquire permits for the entire batch
            # This ensures we don't exceed concurrency limits
            if self.config.has_limits():
                with self.controller:
                    try:
                        batch_results = batch_fn(batch)
                        
                        # Validate batch results
                        if len(batch_results) != len(batch):
                            logger.error(
                                f"Batch {batch_num}: Expected {len(batch)} results, "
                                f"got {len(batch_results)}. Padding with None."
                            )
                            # Pad or truncate to match expected length
                            if len(batch_results) < len(batch):
                                batch_results.extend([None] * (len(batch) - len(batch_results)))
                            else:
                                batch_results = batch_results[:len(batch)]
                        
                        results.extend(batch_results)
                        logger.debug(f"Processed batch {batch_num}/{num_batches}")
                    except Exception as e:
                        logger.error(f"Batch processing error on batch {batch_num}: {e}", exc_info=True)
                        results.extend([None] * len(batch))
            else:
                # No limits, process directly
                try:
                    batch_results = batch_fn(batch)
                    
                    # Validate batch results
                    if len(batch_results) != len(batch):
                        logger.error(
                            f"Batch {batch_num}: Expected {len(batch)} results, "
                            f"got {len(batch_results)}. Padding with None."
                        )
                        if len(batch_results) < len(batch):
                            batch_results.extend([None] * (len(batch) - len(batch_results)))
                        else:
                            batch_results = batch_results[:len(batch)]
                    
                    results.extend(batch_results)
                except Exception as e:
                    logger.error(f"Batch processing error on batch {batch_num}: {e}", exc_info=True)
                    results.extend([None] * len(batch))
        
        return results


# Convenience functions for simple use cases
def process_with_threads(
    items: List[T],
    process_fn: Callable[[T], R],
    max_parallel: Optional[int] = None,
    throttle: Optional[int] = None
) -> List[Optional[R]]:
    """
    Convenience function to process items with ThreadPool.
    
    Use for I/O-bound tasks (API calls, file I/O, network requests).
    
    Args:
        items: List of items to process
        process_fn: Function to apply to each item
        max_parallel: Maximum concurrent threads (None = unlimited)
        throttle: Maximum operations per second (None = unlimited)
    
    Returns:
        List of results in same order as items
    
    Example:
        >>> results = process_with_threads(
        ...     urls,
        ...     fetch_url,
        ...     max_parallel=4,
        ...     throttle=10
        ... )
    """
    config = ParallelismConfig(max_parallel=max_parallel, throttle=throttle)
    executor = ThreadExecutor(config)
    return executor.process_items(items, process_fn)


def process_with_loop(
    items: List[T],
    process_fn: Callable[[T], R],
    throttle: Optional[int] = None
) -> List[Optional[R]]:
    """
    Convenience function to process items sequentially with rate limiting.
    
    Use for CPU/GPU-bound tasks (OCR, model inference).
    
    Args:
        items: List of items to process
        process_fn: Function to apply to each item
        throttle: Maximum operations per second (None = unlimited)
    
    Returns:
        List of results in same order as items
    
    Example:
        >>> results = process_with_loop(
        ...     images,
        ...     ocr_process,
        ...     throttle=10
        ... )
    """
    config = ParallelismConfig(max_parallel=None, throttle=throttle)
    processor = LoopExecutor(config)
    return processor.process_items(items, process_fn)

