"""Token-Erfassung für den Builtin-Gateway.

Parst die `usage`-Felder aus llama-swap-Antworten (Stream + Non-Stream) und meldet
sie an token_stats. Hält den gateway_proxy-Router dünn und ersetzt die zuvor inline
verstreute, still scheiternde String-Suche durch einen testbaren SSE-Zeilenparser.
"""

import json
import logging

from services.token_stats import increment_tokens

log = logging.getLogger(__name__)


def record_usage(usage: dict | None, model: str) -> None:
    """Ein usage-Objekt verbuchen (no-op bei None/leer)."""
    if not usage:
        return
    prompt = usage.get("prompt_tokens", 0)
    completion = usage.get("completion_tokens", 0)
    if prompt or completion:
        increment_tokens(prompt, completion, model=model)


def record_stream_chunk(chunk: bytes, model: str) -> None:
    """Rohen SSE-Chunk auf `usage` prüfen und Tokens verbuchen. Fehler werden
    geloggt (debug) statt verschluckt — ein defekter Chunk bricht den Stream nicht."""
    if b'"usage"' not in chunk:
        return
    text = chunk.decode("utf-8", errors="ignore")
    for line in text.splitlines():
        if not line.startswith("data:"):
            continue
        data_str = line[5:].strip()
        if not data_str or data_str == "[DONE]":
            continue
        try:
            record_usage(json.loads(data_str).get("usage"), model)
        except json.JSONDecodeError:
            log.debug("gateway stream: usage-Parsing fehlgeschlagen: %s", data_str[:120])
