"""Micro-batching layer for non-streaming inference requests.

Accumulates requests arriving within a configurable time window and dispatches
them together to improve GPU utilization.  Streaming requests are never
batched -- they bypass the collector entirely.

The ``BatchCollector`` does not call backends directly.  Instead, it accepts a
``dispatch_fn`` callback at construction time that the dispatcher provides.
Each collected request is dispatched via ``asyncio.gather`` when the batch
flushes (either because the window expires or ``max_size`` is reached).
"""

from __future__ import annotations

import asyncio
import logging
from dataclasses import dataclass, field
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from collections.abc import Awaitable, Callable

    from llmhosts.proxy.models import UnifiedRequest, UnifiedResponse

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Internal bookkeeping
# ---------------------------------------------------------------------------


@dataclass
class _PendingRequest:
    """A request waiting in the batch queue with its associated future."""

    request: UnifiedRequest
    future: asyncio.Future[UnifiedResponse] = field(init=False)

    def __post_init__(self) -> None:
        self.future = asyncio.get_running_loop().create_future()


# ---------------------------------------------------------------------------
# BatchCollector
# ---------------------------------------------------------------------------


class BatchCollector:
    """Micro-batching layer for non-streaming inference requests.

    Collects requests arriving within a configurable window and
    dispatches them as a batch to improve GPU utilization.

    Parameters
    ----------
    dispatch_fn:
        Async callable that sends a single :class:`UnifiedRequest` to a backend
        and returns a :class:`UnifiedResponse`.  Provided by the dispatcher.
    window_ms:
        Maximum time (in milliseconds) to wait before flushing a partial batch.
    max_size:
        Maximum number of requests per batch.  The batch flushes immediately
        when this count is reached, even if the window has not expired.
    enabled:
        When *False*, :meth:`submit` dispatches each request immediately
        without batching (pass-through mode).
    """

    def __init__(
        self,
        dispatch_fn: Callable[[UnifiedRequest], Awaitable[UnifiedResponse]],
        *,
        window_ms: float = 50.0,
        max_size: int = 16,
        enabled: bool = False,
    ) -> None:
        self._dispatch_fn = dispatch_fn
        self._window_ms = window_ms
        self._max_size = max_size
        self._enabled = enabled

        self._pending: list[_PendingRequest] = []
        self._lock: asyncio.Lock = asyncio.Lock()
        self._timer_handle: asyncio.TimerHandle | None = None
        self._flush_task: asyncio.Task[None] | None = None
        self._closed: bool = False

    # -- public properties ---------------------------------------------------

    @property
    def enabled(self) -> bool:
        """Whether micro-batching is active."""
        return self._enabled

    @property
    def window_ms(self) -> float:
        """Current batch window in milliseconds."""
        return self._window_ms

    @property
    def max_size(self) -> int:
        """Maximum batch size before immediate flush."""
        return self._max_size

    @property
    def pending_count(self) -> int:
        """Number of requests currently waiting in the batch."""
        return len(self._pending)

    # -- lifecycle -----------------------------------------------------------

    async def close(self) -> None:
        """Cancel any pending flush timer and reject queued requests."""
        self._closed = True
        self._cancel_timer()
        # Reject any pending requests
        async with self._lock:
            for pending in self._pending:
                if not pending.future.done():
                    pending.future.set_exception(RuntimeError("BatchCollector closed"))
            self._pending.clear()

    # -- submit --------------------------------------------------------------

    async def submit(self, request: UnifiedRequest) -> UnifiedResponse:
        """Submit a request for batching.  Returns when the batch completes.

        Streaming requests (``request.stream is True``) are dispatched
        immediately and never batched.

        When batching is disabled, the request is dispatched immediately
        without queuing.

        Raises
        ------
        RuntimeError
            If the collector has been closed.
        """
        if self._closed:
            raise RuntimeError("BatchCollector is closed")

        # Streaming requests always bypass batching
        if request.stream:
            return await self._dispatch_fn(request)

        # Pass-through when batching disabled
        if not self._enabled:
            return await self._dispatch_fn(request)

        # Enqueue the request
        pending = _PendingRequest(request=request)

        should_flush = False
        async with self._lock:
            self._pending.append(pending)
            if len(self._pending) >= self._max_size:
                should_flush = True

        if should_flush:
            # Drain and flush synchronously (in terms of control flow)
            await self._drain_and_flush()
        else:
            self._ensure_timer()

        return await pending.future

    # -- flush logic ---------------------------------------------------------

    def _cancel_timer(self) -> None:
        """Cancel the window timer if running."""
        if self._timer_handle is not None:
            self._timer_handle.cancel()
            self._timer_handle = None

    def _ensure_timer(self) -> None:
        """Start the window timer if not already running."""
        if self._timer_handle is None:
            loop = asyncio.get_running_loop()
            self._timer_handle = loop.call_later(
                self._window_ms / 1000.0,
                self._on_timer_expired,
            )

    def _on_timer_expired(self) -> None:
        """Callback when the window timer fires.  Schedules a flush."""
        self._timer_handle = None
        self._flush_task = asyncio.create_task(self._drain_and_flush())

    async def _drain_and_flush(self) -> None:
        """Drain the pending queue and dispatch the batch."""
        self._cancel_timer()

        async with self._lock:
            if not self._pending:
                return
            batch = list(self._pending)
            self._pending.clear()

        await self._flush_batch(batch)

    async def _flush_batch(self, batch: list[_PendingRequest]) -> None:
        """Dispatch a collected batch to the backend.

        Each request in the batch is dispatched concurrently via
        ``asyncio.gather``.  Errors are propagated to the individual
        request futures so that one failing request does not affect others.
        """
        if not batch:
            return

        logger.info("Flushing batch of %d request(s)", len(batch))

        async def _dispatch_one(pending: _PendingRequest) -> None:
            try:
                response = await self._dispatch_fn(pending.request)
                if not pending.future.done():
                    pending.future.set_result(response)
            except Exception as exc:
                if not pending.future.done():
                    pending.future.set_exception(exc)

        await asyncio.gather(*[_dispatch_one(p) for p in batch], return_exceptions=True)

    # -- repr ----------------------------------------------------------------

    def __repr__(self) -> str:
        return (
            f"BatchCollector(enabled={self._enabled}, window_ms={self._window_ms}, "
            f"max_size={self._max_size}, pending={len(self._pending)})"
        )
