"""Map local char counts onto real prompt tokens, per model.

``ContextCompressor.estimate_context_tokens`` (chars // 3) is the only token
count available *before* an LLM call, and on this workload it is off by a
factor of ~1.8. Prod sessions measured 101 134 real prompt tokens against a
~56 400 estimate (tool output, JSON schemas and Russian text are token-dense),
so every window decision taken on the estimate — the compression target, the
hard-truncate budgets, the pre-call size check — was taken in a unit roughly
half the size of the real one. Compression aimed at "65% of the window" and
landed above it, the next two tool results re-triggered it, and a session could
run 48 passes in a day.

Every LLM call reports ``prompt_eval_count``, and the char count of the exact
prompt that produced it is known too. Two observations of different size pin a
line::

    real_tokens ≈ overhead + chars / chars_per_token

The line predicts the real cost of any hypothetical prompt assembled from the
same pieces (system text + tool schemas + messages), whichever part of it
changes — a smaller kept region after compression included. Observations are
kept per model: tokenizers differ enough that one line does not cover glm and
gemma.

The fit needs size spread to separate overhead from slope; until it has that,
``DEFAULT_CHARS_PER_TOKEN`` is used. It is deliberately more conservative than
chars // 3 (which was measured at ~1.8x below reality) but not as aggressive as
the worst observed ratio.
"""

from __future__ import annotations

from dataclasses import dataclass

# Fallback until a fitted line exists. The old estimator used 3 chars/token;
# prod measurements sat at 1.67–1.79. 1.8 is the safe middle: closer to truth
# on message text without assuming the token-densest content everywhere.
DEFAULT_CHARS_PER_TOKEN = 1.8

_MAX_OBSERVATIONS = 8
# The two observations must differ this much in size, or the fitted slope is
# noise over a narrow band.
_MIN_SPREAD = 1.25
# Sanity clamps on the fitted slope (tokens per char). A fit outside this band
# means the observations are not comparable (e.g. one prompt was truncated by
# the server) — keep the default instead.
_MIN_TOKENS_PER_CHAR = 1 / 4.0
_MAX_TOKENS_PER_CHAR = 1 / 1.0


@dataclass(frozen=True)
class _Observation:
    chars: int
    tokens: int


class TokenCalibration:
    """Linear fit of prompt tokens against prompt chars for one model."""

    def __init__(self) -> None:
        self._observations: list[_Observation] = []

    def observe(self, chars: int, tokens: int | None) -> None:
        """Record one (prompt chars, real prompt tokens) pair."""
        if not chars or not tokens or tokens <= 0:
            return
        self._observations.append(_Observation(int(chars), int(tokens)))
        del self._observations[: -_MAX_OBSERVATIONS]

    def _line(self) -> tuple[float, float] | None:
        """(overhead, tokens_per_char) or None when the observations can't pin one."""
        n = len(self._observations)
        if n < 2:
            return None
        xs = [o.chars for o in self._observations]
        if max(xs) < min(xs) * _MIN_SPREAD:
            return None

        mean_x = sum(xs) / n
        mean_y = sum(o.tokens for o in self._observations) / n
        sxx = sum((x - mean_x) ** 2 for x in xs)
        if sxx <= 0:
            return None
        sxy = sum((o.chars - mean_x) * (o.tokens - mean_y) for o in self._observations)
        slope = sxy / sxx
        if not (_MIN_TOKENS_PER_CHAR <= slope <= _MAX_TOKENS_PER_CHAR):
            return None
        intercept = mean_y - slope * mean_x
        return max(0.0, intercept), slope

    def tokens_for(self, chars: int) -> int:
        """Estimated real prompt tokens for a prompt of ``chars`` characters."""
        line = self._line()
        if line is None:
            return int(chars / DEFAULT_CHARS_PER_TOKEN)
        intercept, slope = line
        return max(0, round(intercept + chars * slope))

    def delta_for(self, chars: int) -> int:
        """Price *appended* content with the slope only.

        The fitted intercept is the fixed prompt overhead (system prompt and
        tool schemas). That overhead is already inside a real baseline count,
        so charging it again to the delta added on top of that count would
        double-count it.
        """
        line = self._line()
        if line is None:
            return int(chars / DEFAULT_CHARS_PER_TOKEN)
        return max(0, round(chars * line[1]))

    @property
    def observations(self) -> int:
        return len(self._observations)


_registry: dict[str, TokenCalibration] = {}


def _key(model: str | list[str] | None) -> str:
    if isinstance(model, (list, tuple)):
        model = model[0] if model else None
    return (model or "").strip().lower()


def calibration_for(model: str | list[str] | None = None) -> TokenCalibration:
    key = _key(model)
    calibration = _registry.get(key)
    if calibration is None:
        calibration = _registry[key] = TokenCalibration()
    return calibration


def observe_prompt(
    model: str | list[str] | None, chars: int, tokens: int | None
) -> None:
    """Record a real prompt-token count for ``model``."""
    calibration_for(model).observe(chars, tokens)


def reset_calibration() -> None:
    """Drop every observation — tests, and any operator-driven reset."""
    _registry.clear()