"""Tests for token-usage extraction from LangChain-style messages."""

from typing import Any

from skill_eval_runner.usage import message_usage, sum_usage


class _Msg:
    """A message carrying ``usage_metadata`` and/or ``response_metadata``."""

    def __init__(
        self,
        usage_metadata: dict[str, int] | None = None,
        response_metadata: dict[str, Any] | None = None,
    ) -> None:
        if usage_metadata is not None:
            self.usage_metadata = usage_metadata
        self.response_metadata = response_metadata or {}


def test_message_usage_from_usage_metadata() -> None:
    """Reads input/output tokens directly off usage_metadata."""
    msg = _Msg(usage_metadata={"input_tokens": 10, "output_tokens": 4})
    assert message_usage(msg) == (10, 4)


def test_message_usage_falls_back_to_response_metadata() -> None:
    """Falls back to response_metadata['usage'] with prompt/completion keys."""
    msg = _Msg(
        response_metadata={"usage": {"prompt_tokens": 7, "completion_tokens": 2}}
    )
    assert message_usage(msg) == (7, 2)


def test_message_usage_absent_is_zero() -> None:
    """A message with no usage metadata reports zero tokens, not an error."""
    assert message_usage(object()) == (0, 0)


def test_sum_usage_across_messages() -> None:
    """Sums usage across mixed message shapes, skipping those with none."""
    messages = [
        _Msg(usage_metadata={"input_tokens": 10, "output_tokens": 4}),
        _Msg(response_metadata={"usage": {"prompt_tokens": 5, "completion_tokens": 1}}),
        object(),
    ]
    assert sum_usage(messages) == (15, 5)