"""How many requests Pablan lets an endpoint see at once.

A self-hosted llama.cpp server has a fixed number of parallel slots. Sending
more than that does not make it faster: the extra requests sit in the server's
own queue where Pablan can neither see nor bound them, and every one of them
counts against the HTTP timeout. Two colleagues chatting while a reindex runs
is enough to turn a working instance into one where everything times out at
once.

So the waiting happens here instead, in front of the endpoint:

- **One gate per endpoint, not per role.** chat and utility usually point at
  the same server (they do in the shipped `.env`), and it is the SERVER that
  has the slots. Keying by base_url is what makes the limit real.
- **A bounded wait.** A caller waits at most `llm_queue_wait_seconds` for a
  slot and then fails as `llm_busy` — a fast, honest "try again" instead of a
  two-minute timeout that looks like a broken endpoint.
- **A bounded queue.** Past `llm_max_queued` waiters the gate stops admitting:
  when far more work has arrived than the endpoint can absorb, the useful
  answer is "busy", given immediately, to everyone beyond the line.

`llm_busy` is already the vocabulary for this (`llm/errors.py`), and the
frontend phrases it as "the model is busy" — so a queue rejection reaches the
user as the same, correct sentence as a 429 from a cloud provider.

Admin diagnostics (`probe`, `list_models`) deliberately do NOT pass through
the gate: they are single tiny calls, and an admin has to be able to test an
endpoint precisely when it is saturated.
"""

import asyncio
import logging
import time
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager

from app.config import get_settings
from app.llm.errors import LLMError
from app.metrics import metrics

logger = logging.getLogger("pablan.llm")


class _Endpoint:
    """The live picture of one endpoint: who is in it, who is waiting."""

    def __init__(self, limit: int) -> None:
        self.limit = limit
        self.semaphore = asyncio.Semaphore(limit)
        self.in_flight = 0
        self.waiting = 0

    @property
    def saturated(self) -> bool:
        return self.in_flight >= self.limit


_endpoints: dict[str, _Endpoint] = {}


def _busy_error(role: str, reason: str, waited: float) -> LLMError:
    """A queue rejection, in the same shape as an endpoint's own 503 — the
    caller classifies it through `LLMError.code` like any other failure."""
    metrics.inc("llm_queue_rejected_total", {"role": role, "reason": reason})
    logger.info(
        "llm request not admitted",
        extra={
            "event": "llm_queue_rejected",
            "role": role,
            "reason": reason,
            "waited_ms": round(waited * 1000),
        },
    )
    return LLMError(
        f"endpoint busy (role={role}): {reason}",
        role=role,
        kind="queue",
        status="error",
        # 503 is what a saturated endpoint says itself, and what maps to
        # `llm_busy`. Keeping the queue's own rejection in that vocabulary
        # means one reason reaches the user, not two.
        status_code=503,
    )


def _endpoint_for(base_url: str) -> _Endpoint:
    limit = max(1, get_settings().llm_max_parallel)
    endpoint = _endpoints.get(base_url)
    if endpoint is None or endpoint.limit != limit:
        # A changed limit (settings reloaded in a test) rebuilds the gate.
        # In-flight callers hold the old semaphore and still release it.
        endpoint = _Endpoint(limit)
        _endpoints[base_url] = endpoint
    return endpoint


def endpoint_busy(base_url: str) -> bool:
    """Is every slot on this endpoint taken right now?

    Read by the query mode so a waiting turn can SAY it is waiting instead of
    showing a frozen cursor. Advisory: by the time the caller acquires, a slot
    may well have freed.
    """
    endpoint = _endpoints.get(base_url)
    return endpoint is not None and endpoint.saturated


@asynccontextmanager
async def slot(base_url: str, role: str) -> AsyncIterator[None]:
    """Hold one of the endpoint's slots for the whole call.

    For a stream that means until the last token: a streaming completion
    occupies its server slot until it ends, and releasing early would let the
    gate admit work the endpoint has no room for.
    """
    settings = get_settings()
    endpoint = _endpoint_for(base_url)

    if endpoint.saturated and endpoint.waiting >= max(0, settings.llm_max_queued):
        raise _busy_error(role, "queue_full", 0.0)

    started = time.monotonic()
    endpoint.waiting += 1
    metrics.set_gauge("llm_queue_waiting", float(endpoint.waiting), {"role": role})
    try:
        await asyncio.wait_for(
            endpoint.semaphore.acquire(), timeout=settings.llm_queue_wait_seconds
        )
    except TimeoutError:
        raise _busy_error(role, "queue_timeout", time.monotonic() - started) from None
    finally:
        endpoint.waiting -= 1
        metrics.set_gauge("llm_queue_waiting", float(endpoint.waiting), {"role": role})

    waited = time.monotonic() - started
    if waited > 0.01:
        metrics.observe("llm_queue_wait_seconds", waited, {"role": role})
    endpoint.in_flight += 1
    metrics.set_gauge("llm_in_flight", float(endpoint.in_flight), {"role": role})
    try:
        yield
    finally:
        endpoint.in_flight -= 1
        metrics.set_gauge("llm_in_flight", float(endpoint.in_flight), {"role": role})
        endpoint.semaphore.release()


def reset() -> None:
    """Drop every gate. Tests only — a live gate holds waiters."""
    _endpoints.clear()