"""Typed think decode scaffold (Loop 2 P0/P1). No weights.

Gateway-safe: this module must not import train/eval stacks.

P0: prefill ``Answer: <think>\\n``.
P1: allow a real working budget (192 tokens, same as GRPO ``max_new``).
Force-close if still open. Cut only on a triple 20-char repeat (loop XOR),
not on the first equation.
"""

from __future__ import annotations

import re

THINK_CUE = "Show brief working."
THINK_PREFILL = "Answer: <think>\n"
THINK_CLOSE = "</think>"
MAX_THINK_BODY_TOKENS = 192

_LOOP_RE = re.compile(r"(.{20,}?)\1\1", re.I | re.DOTALL)


def whitespace_tokens(text: str) -> list[str]:
    return [t for t in (text or "").split() if t]


def trim_think_body(body: str, *, max_tokens: int = MAX_THINK_BODY_TOKENS) -> str:
    """Keep multi-step work. Cut a triple-repeat loop, then cap at max_tokens."""
    raw = (body or "").replace("<|endoftext|>", "").strip()
    if not raw:
        return ""
    m = _LOOP_RE.search(raw)
    if m:
        raw = raw[: m.start() + len(m.group(1))].strip()
    toks = whitespace_tokens(raw)
    if len(toks) > max_tokens:
        return " ".join(toks[:max_tokens])
    return raw