"""Dynamic batching utilities for face detection inference."""

import logging
import time
from collections import deque
from typing import Callable, Deque, Generic, List, Optional, TypeVar

from gpu_manager import GPUMemoryManager

logger = logging.getLogger(__name__)

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


class DynamicBatcher(Generic[T, R]):
    """Accumulate items into batches and flush based on size or timeout."""

    def __init__(
        self,
        process_fn: Callable[[List[T]], List[R]],
        gpu_manager: Optional[GPUMemoryManager] = None,
        max_batch_size: int = 16,
        batch_timeout_ms: float = 100.0,
        min_batch_size: int = 1,
    ):
        self.process_fn = process_fn
        self.gpu_manager = gpu_manager
        self.max_batch_size = max_batch_size
        self.batch_timeout_ms = batch_timeout_ms
        self.min_batch_size = min_batch_size
        self._queue: Deque[T] = deque()
        self._last_flush = time.monotonic()

    def add(self, item: T) -> List[R]:
        """Add an item and return any results if a batch was flushed."""
        self._queue.append(item)
        if len(self._queue) >= self._current_batch_size():
            return self.flush()

        elapsed_ms = (time.monotonic() - self._last_flush) * 1000
        if elapsed_ms >= self.batch_timeout_ms and len(self._queue) >= self.min_batch_size:
            return self.flush()
        return []

    def flush(self) -> List[R]:
        """Flush all queued items through the processor."""
        if not self._queue:
            return []

        batch = [self._queue.popleft() for _ in range(min(len(self._queue), self._current_batch_size()))]
        results = self.process_fn(batch)
        self._last_flush = time.monotonic()

        if self.gpu_manager is not None:
            self.gpu_manager.adjust_batch_size()
            self.gpu_manager.empty_cache()

        return results

    def _current_batch_size(self) -> int:
        if self.gpu_manager is not None:
            return min(self.gpu_manager.current_batch_size, self.max_batch_size)
        return self.max_batch_size

    @property
    def queued_count(self) -> int:
        return len(self._queue)