"""Reusable numpy array pool to minimize allocation churn during fast image processing."""

from collections import defaultdict
from typing import Dict, List, Tuple

import numpy as np

from wled_controller.utils import get_logger

logger = get_logger(__name__)


class ImagePool:
    """Pool of pre-allocated numpy arrays keyed by shape.

    When a filter needs a new array (e.g. for downscaling), it acquires one
    from the pool instead of allocating. After use, arrays are released back.
    """

    def __init__(self, max_per_shape: int = 4):
        self._pool: Dict[Tuple[int, ...], List[np.ndarray]] = defaultdict(list)
        self._max_per_shape = max_per_shape

    def acquire(self, height: int, width: int, channels: int = 3) -> np.ndarray:
        """Get a pre-allocated array or create a new one."""
        shape = (height, width, channels)
        bucket = self._pool[shape]
        if bucket:
            return bucket.pop()
        return np.empty(shape, dtype=np.uint8)

    def release(self, array: np.ndarray) -> None:
        """Return an array to the pool for reuse."""
        shape = array.shape
        bucket = self._pool[shape]
        if len(bucket) < self._max_per_shape:
            bucket.append(array)

    def clear(self) -> None:
        """Release all pooled arrays."""
        self._pool.clear()
