"""
Общий модуль извлечения token usage из ответа OmniRoute.

Используется и cron-пробой (probe.py), и ручной проверкой (main.py).
Правила:
  * usage извлекается ДО валидации контента — чтобы провал content-парсинга
    не терял реально потраченные токены;
  * если провайдер не отдал usage (или отдал полностью нулевой) — все поля
    остаются NULL: неизвестное потребление не подменяется нулями;
  * usage_json хранит сырой словарь usage как пришел от провайдера.
"""
from __future__ import annotations

import json
from typing import Optional

# Возможные имена полей в разных диалектах (OpenAI / Anthropic / camelCase).
_INPUT_KEYS = ("prompt_tokens", "input_tokens", "promptTokens", "inputTokens")
_OUTPUT_KEYS = ("completion_tokens", "output_tokens", "completionTokens", "outputTokens")
_TOTAL_KEYS = ("total_tokens", "totalTokens")


def _first_int(usage: dict, keys: tuple[str, ...]) -> Optional[int]:
    """Первое числовое значение из ключей; не-числа и None игнорируются."""
    for k in keys:
        v = usage.get(k)
        if isinstance(v, bool):  # bool — подкласс int, отсекаем явно
            continue
        if isinstance(v, (int, float)) and v >= 0:
            return int(v)
    return None


def extract_usage(resp) -> dict:
    """Достать usage из тела ответа chat/completions.

    Возвращает dict с ключами input_tokens / output_tokens / total_tokens /
    usage_json. Отсутствующий или вырожденный (полностью нулевой) usage дает
    None во всех полях — NULL, а не нулевая оценка.
    """
    empty = {"input_tokens": None, "output_tokens": None,
             "total_tokens": None, "usage_json": None}

    usage = None
    if isinstance(resp, dict):
        u = resp.get("usage")
        if isinstance(u, dict):
            usage = u
        elif isinstance(u, str):
            # редко: usage пришел строкой-JSON
            try:
                parsed = json.loads(u)
                if isinstance(parsed, dict):
                    usage = parsed
            except (ValueError, TypeError):
                usage = None
    if not isinstance(usage, dict):
        return empty

    inp = _first_int(usage, _INPUT_KEYS)
    out = _first_int(usage, _OUTPUT_KEYS)
    tot = _first_int(usage, _TOTAL_KEYS)

    # Полностью нулевой usage (0/0/0) трактуем как «не отдал» — провайдеры
    # так отвечают в ошибочных путях; реальные токены всегда > 0.
    if not any(v for v in (inp, out, tot) if v is not None):
        return empty

    # total — точная арифметика по отчетным значениям, не оценка.
    if tot is None and inp is not None and out is not None:
        tot = inp + out

    try:
        usage_json = json.dumps(usage, ensure_ascii=False, sort_keys=True)
    except (TypeError, ValueError):
        usage_json = None

    return {
        "input_tokens": inp,
        "output_tokens": out,
        "total_tokens": tot,
        "usage_json": usage_json,
    }