"""Edge-case tests for metrics/token_counter.py — usage metadata, overflow, and error tracking."""

from unittest.mock import Mock


def _make_callback(**kwargs):
    """Create a TokenCountingCallback without database dependencies."""
    from local_deep_research.metrics.token_counter import TokenCountingCallback

    return TokenCountingCallback(**kwargs)


class TestUsageMetadataExtraction:
    """Tests for the usage_metadata extraction path (lines 248-262)."""

    def test_usage_metadata_extraction_path(self):
        """usage_metadata path used by Google/Gemini providers."""
        callback = _make_callback()
        callback.current_model = "gemini-pro"
        callback.current_provider = "google"
        callback.counts["by_model"]["gemini-pro"] = {
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "total_tokens": 0,
            "calls": 1,
            "provider": "google",
        }

        # Build a mock LLMResult with usage_metadata on the message
        response = Mock()
        response.llm_output = None

        message = Mock()
        message.usage_metadata = {
            "input_tokens": 100,
            "output_tokens": 50,
            "total_tokens": 150,
        }
        message.response_metadata = {}

        generation = Mock()
        generation.message = message

        response.generations = [[generation]]

        callback.on_llm_end(response)

        assert callback.counts["total_prompt_tokens"] == 100
        assert callback.counts["total_completion_tokens"] == 50
        assert callback.counts["total_tokens"] == 150


class TestContextOverflowDetection:
    """Tests for context overflow edge cases."""

    def test_no_overflow_without_context_limit(self):
        """context_limit=None must NOT trigger false overflow warnings."""
        callback = _make_callback()
        callback.context_limit = None
        callback.current_model = "test-model"
        callback.current_provider = "test"
        callback.counts["by_model"]["test-model"] = {
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "total_tokens": 0,
            "calls": 1,
            "provider": "test",
        }

        # Build response with high prompt_eval_count but no context_limit
        response = Mock()
        response.llm_output = None

        message = Mock()
        message.usage_metadata = None
        message.response_metadata = {
            "prompt_eval_count": 50000,
            "eval_count": 1000,
        }
        generation = Mock()
        generation.message = message
        response.generations = [[generation]]

        callback.on_llm_end(response)

        # No overflow should be detected when context_limit is None
        assert callback.context_truncated is False

    def test_overflow_with_empty_prompts_no_divzero(self):
        """original_prompt_estimate=0 doesn't cause division-by-zero in truncation_ratio."""
        callback = _make_callback()
        callback.context_limit = 4096
        callback.original_prompt_estimate = 0
        callback.current_model = "test-model"
        callback.current_provider = "test"
        callback.counts["by_model"]["test-model"] = {
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "total_tokens": 0,
            "calls": 1,
            "provider": "test",
        }

        response = Mock()
        response.llm_output = None

        message = Mock()
        message.usage_metadata = None
        message.response_metadata = {
            "prompt_eval_count": 4000,
            "eval_count": 100,
        }
        generation = Mock()
        generation.message = message
        response.generations = [[generation]]

        # Should not raise ZeroDivisionError
        callback.on_llm_end(response)
        assert callback.context_truncated is True
        # truncation_ratio should be 0 since original_prompt_estimate is 0
        assert callback.truncation_ratio == 0

    def test_overflow_fields_none_when_no_overflow(self):
        """_get_context_overflow_fields() returns None for tokens_truncated when no overflow."""
        callback = _make_callback()
        callback.context_truncated = False
        callback.tokens_truncated = 0

        fields = callback._get_context_overflow_fields()
        assert fields["tokens_truncated"] is None
        assert fields["truncation_ratio"] is None
        assert fields["context_truncated"] is False


class TestTokenCountAccumulation:
    """Tests for in-memory count accumulation."""

    def test_counts_accumulate_across_calls(self):
        """In-memory counts correctly sum over 3+ LLM calls."""
        callback = _make_callback()
        callback.current_model = "test-model"
        callback.current_provider = "test"
        callback.counts["by_model"]["test-model"] = {
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "total_tokens": 0,
            "calls": 0,
            "provider": "test",
        }

        def make_response(prompt_tokens, completion_tokens):
            resp = Mock()
            resp.llm_output = {
                "token_usage": {
                    "prompt_tokens": prompt_tokens,
                    "completion_tokens": completion_tokens,
                    "total_tokens": prompt_tokens + completion_tokens,
                }
            }
            return resp

        callback.on_llm_end(make_response(10, 5))
        callback.on_llm_end(make_response(20, 10))
        callback.on_llm_end(make_response(30, 15))

        assert callback.counts["total_prompt_tokens"] == 60
        assert callback.counts["total_completion_tokens"] == 30
        assert callback.counts["total_tokens"] == 90
        assert callback.counts["by_model"]["test-model"]["prompt_tokens"] == 60


class TestOnLlmError:
    """Tests for on_llm_error tracking."""

    def test_on_llm_error_captures_type_and_timing(self):
        """Error tracking captures error_type='ValueError' and response_time_ms."""
        import time

        callback = _make_callback()
        callback.start_time = time.time() - 0.5  # 500ms ago

        error = ValueError("Something went wrong")
        callback.on_llm_error(error)

        assert callback.success_status == "error"
        assert callback.error_type == "ValueError"
        assert callback.response_time_ms is not None
        assert callback.response_time_ms >= 400  # At least ~400ms


class TestOllamaOverflowDetection:
    """Tests for Ollama-specific overflow detection."""

    def test_response_metadata_ollama_overflow_detection(self):
        """prompt_eval_count >= 95% of context_limit triggers context_truncated=True."""
        callback = _make_callback()
        callback.context_limit = 4096
        callback.current_model = "llama3"
        callback.current_provider = "ollama"
        callback.original_prompt_estimate = 5000
        callback.counts["by_model"]["llama3"] = {
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "total_tokens": 0,
            "calls": 1,
            "provider": "ollama",
        }

        response = Mock()
        response.llm_output = None

        message = Mock()
        message.usage_metadata = None
        message.response_metadata = {
            "prompt_eval_count": 3900,  # 95.2% of 4096
            "eval_count": 200,
        }
        generation = Mock()
        generation.message = message
        response.generations = [[generation]]

        callback.on_llm_end(response)

        assert callback.context_truncated is True
        assert callback.tokens_truncated > 0
