"""Intelligent batching for Azure API calls."""

import asyncio
import re
from typing import TypeVar, List, Callable, AsyncIterator, Optional

from ..utils.logging import get_logger
from .rate_control import AdaptiveRateController

LOGGER = get_logger()

T = TypeVar("T")
R = TypeVar("R")


class BatchProcessor:
    """Process items in adaptive batch sizes with concurrency control."""

    def __init__(
        self,
        initial_batch_size: int = 100,
        min_batch_size: int = 10,
        max_batch_size: int = 1000,
        max_concurrent: int = 5,
    ):
        self.current_batch_size = initial_batch_size
        self.min_batch_size = min_batch_size
        self.max_batch_size = max_batch_size
        self.max_concurrent = max_concurrent
        self.semaphore = asyncio.Semaphore(max_concurrent)

    async def process_batches(
        self,
        items: List[T],
        processor: Callable[[List[T]], R],
        rate_controller: Optional[AdaptiveRateController] = None,
    ) -> AsyncIterator[R]:
        """Process items in adaptive batches with concurrency control.

        Args:
            items: Items to process
            processor: Async function to process a batch
            rate_controller: Optional rate controller for throttling

        Yields:
            Results from each batch
        """
        idx = 0
        tasks = []

        while idx < len(items):
            batch = items[idx : idx + self.current_batch_size]
            idx += self.current_batch_size

            if rate_controller:
                await rate_controller.acquire()

            task = asyncio.create_task(
                self._process_batch_with_retry(batch, processor, rate_controller)
            )
            tasks.append(task)

            # Limit concurrent batches
            if len(tasks) >= self.max_concurrent:
                done, pending = await asyncio.wait(
                    tasks, return_when=asyncio.FIRST_COMPLETED
                )
                tasks = list(pending)
                for completed_task in done:
                    try:
                        result = await completed_task
                        yield result
                    except Exception as exc:
                        LOGGER.error(
                            "Batch processing failed",
                            extra={"context": {"error": str(exc)}},
                        )

        # Wait for remaining tasks
        if tasks:
            results = await asyncio.gather(*tasks, return_exceptions=True)
            for result in results:
                if not isinstance(result, Exception):
                    yield result

    async def _process_batch_with_retry(
        self,
        batch: List[T],
        processor: Callable[[List[T]], R],
        rate_controller: Optional[AdaptiveRateController] = None,
        attempt: int = 0,
    ) -> R:
        """Process batch with automatic retry and backoff.

        The semaphore is acquired per-attempt and released before any
        retry sleep so that other batches are not blocked during backoff.

        Args:
            batch: Batch of items to process
            processor: Processing function
            rate_controller: Optional rate controller
            attempt: Current retry attempt

        Returns:
            Processing result

        Raises:
            Exception: If all retries exhausted
        """
        last_exc: Optional[Exception] = None
        max_retries = 3

        for current_attempt in range(attempt, attempt + max_retries + 1):
            async with self.semaphore:
                try:
                    result = await processor(batch)

                    if rate_controller:
                        await rate_controller.record_success()

                    # Increase batch size on success
                    if self.current_batch_size < self.max_batch_size:
                        self.current_batch_size = min(
                            self.max_batch_size, int(self.current_batch_size * 1.5)
                        )

                    return result

                except Exception as exc:
                    last_exc = exc
                    if not self._is_throttle_error(exc):
                        raise

                    if rate_controller:
                        retry_after = self._extract_retry_after(exc)
                        await rate_controller.record_throttle(retry_after)

                    # Decrease batch size
                    self.current_batch_size = max(
                        self.min_batch_size, int(self.current_batch_size * 0.5)
                    )

            # Semaphore released — sleep before retrying
            if current_attempt < attempt + max_retries:
                delay = 2**current_attempt
                LOGGER.info(
                    f"Throttled, retrying batch after {delay}s",
                    extra={
                        "context": {
                            "attempt": current_attempt + 1,
                            "batch_size": len(batch),
                            "new_batch_size": self.current_batch_size,
                        }
                    },
                )
                await asyncio.sleep(delay)

        # All retries exhausted
        raise last_exc or Exception("Batch processing failed after retries")

    def _is_throttle_error(self, exception: Exception) -> bool:
        """Check if exception indicates API throttling.

        Args:
            exception: Exception to check

        Returns:
            True if throttling detected
        """
        error_str = str(exception).lower()
        return (
            "429" in error_str
            or "throttl" in error_str
            or "rate limit" in error_str
            or "too many requests" in error_str
        )

    def _extract_retry_after(self, exception: Exception) -> Optional[int]:
        """Extract Retry-After value from exception.

        Args:
            exception: Exception potentially containing retry info

        Returns:
            Retry delay in seconds, or None if not found
        """
        # Try to extract from exception message
        match = re.search(r"retry[_\s-]?after[:\s]+(\d+)", str(exception), re.I)
        if match:
            return int(match.group(1))

        # Check exception attributes
        if hasattr(exception, "retry_after"):
            return int(exception.retry_after)

        return None

    def get_current_batch_size(self) -> int:
        """Get current batch size.

        Returns:
            Current batch size setting
        """
        return self.current_batch_size
