"""Utilities for normalizing LLM usage metadata."""
from typing import Dict, Any


def normalize_usage(usage: Dict[str, Any]) -> Dict[str, int]:
    """Normalize usage keys to input/output/total integers.

    Supports variants like prompt_tokens/completion_tokens while being defensive.
    """
    try:
        input_val = usage.get("input_tokens")
        if not isinstance(input_val, (int, float)):
            input_val = usage.get("prompt_tokens", 0)

        output_val = usage.get("output_tokens")
        if not isinstance(output_val, (int, float)):
            output_val = usage.get("completion_tokens", 0)

        total_val = usage.get("total_tokens")
        if not isinstance(total_val, (int, float)):
            total_val = (int(input_val or 0)) + (int(output_val or 0))

        return {
            "input_tokens": int(input_val or 0),
            "output_tokens": int(output_val or 0),
            "total_tokens": int(total_val or 0),
        }
    except Exception:
        return {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0}