"""Tests for the per-run token usage accumulator and its context-scoped activation."""
from __future__ import annotations

from rag_eval.metrics.token_tracker import (
    TokenUsageTracker,
    get_current_tracker,
    track_token_usage,
)


def test_record_accumulates_input_output_and_calls():
    tracker = TokenUsageTracker()
    tracker.record("gpt-5", 100, 50)
    tracker.record("gpt-5", 20, 10)
    assert tracker.as_dict() == {
        "gpt-5": {"input_tokens": 120, "output_tokens": 60, "calls": 2}
    }


def test_record_groups_by_model_name():
    tracker = TokenUsageTracker()
    tracker.record("gpt-5", 100, 50)
    tracker.record("Qwen/Qwen3-Embedding-4B", 30, 0)
    result = tracker.as_dict()
    assert set(result.keys()) == {"gpt-5", "Qwen/Qwen3-Embedding-4B"}
    assert result["Qwen/Qwen3-Embedding-4B"] == {
        "input_tokens": 30, "output_tokens": 0, "calls": 1
    }


def test_record_defaults_blank_model_to_unknown():
    tracker = TokenUsageTracker()
    tracker.record("", 10, 5)
    assert "unknown" in tracker.as_dict()


def test_merge_into_sums_with_existing_totals():
    tracker = TokenUsageTracker()
    tracker.record("gpt-5", 100, 50)
    existing = {"gpt-5": {"input_tokens": 200, "output_tokens": 100, "calls": 3}}
    merged = tracker.merge_into(existing)
    assert merged == {"gpt-5": {"input_tokens": 300, "output_tokens": 150, "calls": 4}}


def test_merge_into_keeps_models_only_in_existing():
    tracker = TokenUsageTracker()
    tracker.record("gpt-5", 10, 5)
    existing = {"other-model": {"input_tokens": 1, "output_tokens": 1, "calls": 1}}
    merged = tracker.merge_into(existing)
    assert merged["other-model"] == {"input_tokens": 1, "output_tokens": 1, "calls": 1}
    assert merged["gpt-5"] == {"input_tokens": 10, "output_tokens": 5, "calls": 1}


def test_merge_into_does_not_mutate_existing_dict():
    tracker = TokenUsageTracker()
    tracker.record("gpt-5", 10, 5)
    existing = {"gpt-5": {"input_tokens": 1, "output_tokens": 1, "calls": 1}}
    tracker.merge_into(existing)
    assert existing == {"gpt-5": {"input_tokens": 1, "output_tokens": 1, "calls": 1}}


def test_get_current_tracker_returns_none_outside_context():
    assert get_current_tracker() is None


def test_track_token_usage_activates_and_resets_context():
    assert get_current_tracker() is None
    with track_token_usage() as tracker:
        assert get_current_tracker() is tracker
        tracker.record("gpt-5", 1, 1)
    assert get_current_tracker() is None


def test_track_token_usage_nested_contexts_are_isolated():
    with track_token_usage() as outer:
        outer.record("outer-model", 5, 5)
        with track_token_usage() as inner:
            inner.record("inner-model", 1, 1)
            assert get_current_tracker() is inner
        assert get_current_tracker() is outer
    assert outer.as_dict() == {
        "outer-model": {"input_tokens": 5, "output_tokens": 5, "calls": 1}
    }