"""Bounded, ordered prefill transfer primitives.

Prefill chunks mutate the downstream shard's session cache, so they must reach a
route in order.  This deliberately uses a serial acknowledgement window: it is
the safe default for both current peers and old peers which do not advertise a
windowing capability.  The configured in-flight limit is still explicit so a
future ordered transport can widen the window without changing callers.
"""

from __future__ import annotations

import os
from dataclasses import dataclass
from threading import Event
from typing import Callable, Iterable, TypeVar


DEFAULT_PREFILL_CHUNK_TOKENS = 128
DEFAULT_PREFILL_MAX_IN_FLIGHT = 1
DEFAULT_PREFILL_MAX_CHUNK_BYTES = 8 * 1024 * 1024

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


@dataclass(frozen=True)
class PrefillTransferLimits:
    """Configuration for one ordered prefill seam."""

    chunk_tokens: int = DEFAULT_PREFILL_CHUNK_TOKENS
    max_in_flight: int = DEFAULT_PREFILL_MAX_IN_FLIGHT
    max_chunk_bytes: int = DEFAULT_PREFILL_MAX_CHUNK_BYTES

    @property
    def effective_in_flight(self) -> int:
        """Current peers require ordered session-cache mutation, hence one ack."""
        return 1

    @property
    def max_buffered_bytes(self) -> int:
        """Hard accounting bound, including any future wider ack window."""
        return self.max_chunk_bytes * self.max_in_flight

    @classmethod
    def from_env(cls) -> "PrefillTransferLimits":
        # MESHNET_CHUNK_TOKENS was the pre-DIP-007 name.  Keep it as a fallback
        # so existing deployments retain their chunk shape while upgrading.
        return cls(
            chunk_tokens=_positive_env(
                "MESHNET_PREFILL_CHUNK_TOKENS",
                _positive_env("MESHNET_CHUNK_TOKENS", DEFAULT_PREFILL_CHUNK_TOKENS),
            ),
            max_in_flight=_positive_env(
                "MESHNET_PREFILL_MAX_IN_FLIGHT", DEFAULT_PREFILL_MAX_IN_FLIGHT,
            ),
            max_chunk_bytes=_positive_env(
                "MESHNET_PREFILL_MAX_CHUNK_BYTES", DEFAULT_PREFILL_MAX_CHUNK_BYTES,
            ),
        )


class BoundedPrefillSender:
    """Send lazily-produced chunks with bounded ownership and ordered acks."""

    def __init__(self, limits: PrefillTransferLimits) -> None:
        self.limits = limits
        self.buffered_bytes = 0
        self.peak_buffered_bytes = 0
        self.in_flight = 0
        self.peak_in_flight = 0
        self.closed = False

    def send(
        self,
        chunks: Iterable[T],
        *,
        body_size: Callable[[T], int],
        forward: Callable[[T], R],
        cancelled: Event | None = None,
    ) -> list[R]:
        """Forward chunks in source order, releasing each body after its ack.

        ``forward`` is synchronous by design: a slow consumer therefore blocks
        production of the next chunk instead of accumulating an unbounded queue.
        Every retained body is dropped on cancellation or route failure.
        """
        results: list[R] = []
        try:
            for chunk in chunks:
                if self.closed or (cancelled is not None and cancelled.is_set()):
                    break
                size = body_size(chunk)
                if size < 0:
                    raise ValueError("prefill chunk size cannot be negative")
                if size > self.limits.max_chunk_bytes:
                    raise ValueError(
                        f"prefill chunk exceeds {self.limits.max_chunk_bytes} byte limit"
                    )
                self.buffered_bytes += size
                self.in_flight += 1
                self.peak_buffered_bytes = max(self.peak_buffered_bytes, self.buffered_bytes)
                self.peak_in_flight = max(self.peak_in_flight, self.in_flight)
                try:
                    results.append(forward(chunk))
                finally:
                    # Do not retain a body while waiting for the next chunk.
                    self.buffered_bytes -= size
                    self.in_flight -= 1
        except BaseException:
            self.close()
            raise
        return results

    def close(self) -> None:
        """Release accounting after cancellation or route failure.

        The sender deliberately owns no queued chunk references; callers must
        discard their iterator on close rather than trying to drain it.
        """
        self.buffered_bytes = 0
        self.in_flight = 0
        self.closed = True


def _positive_env(name: str, default: int) -> int:
    try:
        value = int(os.environ.get(name, default))
    except (TypeError, ValueError):
        return default
    return value if value > 0 else default