"""Tests for app.callbacks module."""

from __future__ import annotations

import pytest

from app.callbacks import TokenUsageCallbackHandler


@pytest.mark.unit
class TestTokenUsageCallbackHandler:
    def test_initial_state(self) -> None:
        handler = TokenUsageCallbackHandler(model_name="claude-sonnet-4-6")
        usage = handler.get_usage()
        assert usage.prompt_tokens == 0
        assert usage.completion_tokens == 0
        assert usage.total_tokens == 0
        assert usage.total_cost_usd == 0.0

    def test_accumulates_tokens(self) -> None:
        handler = TokenUsageCallbackHandler(model_name="claude-sonnet-4-6")

        class FakeResult:
            llm_output = {"token_usage": {"prompt_tokens": 100, "completion_tokens": 50}}

        handler.on_llm_end(FakeResult())
        usage = handler.get_usage()
        assert usage.prompt_tokens == 100
        assert usage.completion_tokens == 50
        assert usage.total_tokens == 150

    def test_accumulates_across_calls(self) -> None:
        handler = TokenUsageCallbackHandler(model_name="claude-sonnet-4-6")

        class FakeResult:
            llm_output = {"token_usage": {"prompt_tokens": 100, "completion_tokens": 50}}

        handler.on_llm_end(FakeResult())
        handler.on_llm_end(FakeResult())
        usage = handler.get_usage()
        assert usage.prompt_tokens == 200
        assert usage.completion_tokens == 100
        assert usage.total_tokens == 300

    def test_cost_calculation(self) -> None:
        handler = TokenUsageCallbackHandler(model_name="claude-sonnet-4-6")

        class FakeResult:
            llm_output = {"token_usage": {"prompt_tokens": 1000, "completion_tokens": 1000}}

        handler.on_llm_end(FakeResult())
        usage = handler.get_usage()
        # claude-sonnet-4-6: prompt $0.003/1K, completion $0.015/1K
        expected_cost = 1000 * 0.003 / 1000 + 1000 * 0.015 / 1000
        assert usage.total_cost_usd == pytest.approx(expected_cost)

    def test_reset(self) -> None:
        handler = TokenUsageCallbackHandler(model_name="claude-sonnet-4-6")

        class FakeResult:
            llm_output = {"token_usage": {"prompt_tokens": 100, "completion_tokens": 50}}

        handler.on_llm_end(FakeResult())
        handler.reset()
        usage = handler.get_usage()
        assert usage.total_tokens == 0

    def test_usage_is_immutable(self) -> None:
        handler = TokenUsageCallbackHandler()
        usage = handler.get_usage()
        with pytest.raises(Exception):
            usage.prompt_tokens = 999

    def test_unknown_model_uses_default_cost(self) -> None:
        handler = TokenUsageCallbackHandler(model_name="unknown-model")

        class FakeResult:
            llm_output = {"token_usage": {"prompt_tokens": 1000, "completion_tokens": 1000}}

        handler.on_llm_end(FakeResult())
        usage = handler.get_usage()
        assert usage.total_cost_usd > 0

    def test_handles_missing_token_usage(self) -> None:
        handler = TokenUsageCallbackHandler()

        class FakeResult:
            llm_output = {}

        handler.on_llm_end(FakeResult())
        usage = handler.get_usage()
        assert usage.total_tokens == 0

    def test_handles_none_llm_output(self) -> None:
        handler = TokenUsageCallbackHandler()

        class FakeResult:
            llm_output = None

        handler.on_llm_end(FakeResult())
        usage = handler.get_usage()
        assert usage.total_tokens == 0