"""Model-aware wall-clock budgeting.

CPU inference time is U-shaped in model size: tiny models are slow because they
*over-generate* (rambling completions that hit num_predict every call) and retry
more, while big models are slow *per token* (more parameters). 3-4B is the sweet
spot. Reasoning models (r1/qwq/o1) add a long <think> preamble on top.

We scale only the *wall-clock* budget — never the implement-attempt count — so
the correctness and speed axes stay comparable across model sizes (attempts feed
``scoring._speed_score``'s cap; changing them per model would change what the
benchmark measures). Mirrors ``bench-sandbox``'s reasoning-model 3x rule.
"""

from __future__ import annotations

import re

# Same tags bench-sandbox uses, so the whole suite treats reasoning models alike.
_REASONING_TAGS = ("r1", "qwq", "reason", "think", "o1")
_PARAM_RE = re.compile(r"(\d+(?:\.\d+)?)\s*b\b", re.IGNORECASE)


def _is_reasoning(model: str | None) -> bool:
    return bool(model) and any(t in model.lower() for t in _REASONING_TAGS)


def _param_b(model: str | None) -> float | None:
    """Best-effort parameter count in billions parsed from the tag (e.g.
    ``qwen2.5-coder:1.5b`` -> 1.5). Returns None when the tag carries no size."""
    if not model:
        return None
    m = _PARAM_RE.search(model)
    return float(m.group(1)) if m else None


def budget_scale(model: str | None) -> float:
    """Multiplier applied to the base wall-clock budget for ``model``.

    Reasoning dominates (it stacks a think-preamble on whatever size it is).
    Otherwise the U-shaped size curve applies; unknown sizes get the 1.0
    baseline. Defaults are first-guess and meant to be retuned from the ledger's
    recorded ``wall_clock_s`` vs ``wall_clock_budget``."""
    if _is_reasoning(model):
        return 3.0
    b = _param_b(model)
    if b is None:
        return 1.0
    if b <= 1.5:
        return 1.5      # over-generation + retries
    if b <= 4:
        return 1.0      # sweet spot
    if b < 13:
        return 1.5      # per-token CPU latency (7b-class)
    return 2.0          # large models, more so