"""Compatibility-aware microbatch scheduling with bounded lifecycle semantics."""

from __future__ import annotations
import asyncio
import time
from collections import defaultdict
from dataclasses import dataclass
from .backends.base import Backend, BackendInput, BackendResult
from .errors import RequestLimitError


@dataclass
class _Pending:
    item: BackendInput
    submitted_at: float
    future: asyncio.Future[tuple[BackendResult, float, int]]


class BatchScheduler:
    def __init__(
        self,
        backend: Backend,
        *,
        max_batch_size: int = 8,
        batch_window_ms: int = 8,
        max_pending_requests: int = 256,
    ):
        if (
            type(max_batch_size) is not int
            or max_batch_size < 1
            or type(max_pending_requests) is not int
            or max_pending_requests < 1
        ):
            raise ValueError("Batch and queue capacities must be positive integers")
        if type(batch_window_ms) is not int or batch_window_ms < 0:
            raise ValueError("Batch window must be a nonnegative integer")
        self.backend = backend
        self.max_batch_size = max_batch_size
        self.batch_window_ms = batch_window_ms
        self._queue: asyncio.Queue[_Pending | None] = asyncio.Queue(
            maxsize=max_pending_requests
        )
        self._worker: asyncio.Task[None] | None = None
        self._closing_task: asyncio.Task[None] | None = None
        self._closed = False
        self._failure: RuntimeError | None = None

    async def submit(self, item: BackendInput) -> tuple[BackendResult, float, int]:
        if self._closed:
            raise RuntimeError("Batch scheduler is closed")
        if self._failure is not None:
            raise self._failure
        loop = asyncio.get_running_loop()
        if self._worker is None:
            self._worker = loop.create_task(self._run(), name="backplane-batcher")
        future = loop.create_future()
        try:
            self._queue.put_nowait(_Pending(item, time.perf_counter(), future))
        except asyncio.QueueFull as exc:
            future.cancel()
            raise RequestLimitError("Runtime pending-request queue is full") from exc
        return await future

    async def _run(self):
        pending = []
        try:
            while True:
                first = await self._queue.get()
                if first is None:
                    return
                pending = [first]
                deadline = time.perf_counter() + self.batch_window_ms / 1000
                while len(pending) < self.max_batch_size:
                    remaining = deadline - time.perf_counter()
                    if remaining <= 0:
                        break
                    try:
                        candidate = await asyncio.wait_for(self._queue.get(), remaining)
                    except TimeoutError:
                        break
                    if candidate is None:
                        # Submissions are already closed, so the marker is the
                        # last item and can be replaced without blocking.
                        self._queue.put_nowait(None)
                        break
                    pending.append(candidate)
                groups = defaultdict(list)
                for entry in pending:
                    if not entry.future.done():
                        groups[entry.item.stack_key].append(entry)
                for entries in groups.values():
                    entries = [entry for entry in entries if not entry.future.done()]
                    if not entries:
                        continue
                    try:
                        started = time.perf_counter()
                        results = await self.backend.generate_batch(
                            [entry.item for entry in entries]
                        )
                        if len(results) != len(entries):
                            raise RuntimeError(
                                "backend returned the wrong batch cardinality"
                            )
                        for entry, result in zip(entries, results, strict=True):
                            if not entry.future.done():
                                entry.future.set_result(
                                    (
                                        result,
                                        (started - entry.submitted_at) * 1000,
                                        len(entries),
                                    )
                                )
                    except Exception as exc:
                        for entry in entries:
                            if not entry.future.done():
                                entry.future.set_exception(exc)
                pending = []
        except asyncio.CancelledError:
            self._failure = RuntimeError(
                "Batch worker was cancelled before completing its requests"
            )
        except Exception as exc:
            self._failure = RuntimeError("Batch worker stopped unexpectedly")
            self._failure.__cause__ = exc
        finally:
            failure = self._failure or RuntimeError(
                "Batch scheduler stopped before completing the request"
            )
            for entry in pending:
                if not entry.future.done():
                    entry.future.set_exception(failure)
            while not self._queue.empty():
                entry = self._queue.get_nowait()
                if entry is not None and not entry.future.done():
                    entry.future.set_exception(failure)

    async def _shutdown(self):
        if self._worker is not None and not self._worker.done():
            await self._queue.put(None)
            await self._worker

    async def close(self):
        if self._closing_task is None:
            self._closed = True
            self._closing_task = asyncio.create_task(
                self._shutdown(), name="backplane-batcher-close"
            )
        await asyncio.shield(self._closing_task)
