"""
Extended Tests for Token Counter

Phase 17: Token Counter & Metrics - Extended Token Counting Tests
Tests token counting, call stack tracking, and research context integration.
"""

import time
from unittest.mock import MagicMock

from local_deep_research.metrics.token_counter import (
    TokenCountingCallback,
)


class TestTokenCounting:
    """Tests for basic token counting functionality"""

    def test_on_llm_start_token_estimation(self):
        """Test token estimation from prompt length"""
        callback = TokenCountingCallback(research_id="test-123")

        prompts = ["This is a test prompt with about 40 characters"]
        serialized = {"_type": "ChatOpenAI"}

        callback.on_llm_start(serialized, prompts)

        # ~40 chars / 4 = ~10 tokens estimated
        assert callback.original_prompt_estimate > 0
        assert callback.original_prompt_estimate < 20

    def test_on_llm_end_token_extraction(self):
        """Test token extraction from LLM response"""
        callback = TokenCountingCallback(research_id="test-123")

        # Setup start - pass model via invocation_params
        callback.on_llm_start(
            {"_type": "ChatOpenAI"},
            ["test"],
            invocation_params={"model": "gpt-4"},
        )

        # Create mock response with token usage
        mock_response = MagicMock()
        mock_response.llm_output = {
            "token_usage": {
                "prompt_tokens": 100,
                "completion_tokens": 50,
                "total_tokens": 150,
            }
        }
        mock_response.generations = []

        callback.on_llm_end(mock_response)

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

    def test_prompt_tokens_counting(self):
        """Test prompt tokens are counted correctly"""
        callback = TokenCountingCallback()

        callback.on_llm_start(
            {"_type": "ChatOpenAI"},
            ["test"],
            invocation_params={"model": "gpt-4"},
        )

        mock_response = MagicMock()
        mock_response.llm_output = {
            "token_usage": {
                "prompt_tokens": 200,
                "completion_tokens": 0,
                "total_tokens": 200,
            }
        }
        mock_response.generations = []

        callback.on_llm_end(mock_response)

        assert callback.counts["total_prompt_tokens"] == 200
        assert callback.counts["by_model"]["gpt-4"]["prompt_tokens"] == 200

    def test_completion_tokens_counting(self):
        """Test completion tokens are counted correctly"""
        callback = TokenCountingCallback()

        callback.on_llm_start(
            {"_type": "ChatOpenAI"},
            ["test"],
            invocation_params={"model": "gpt-4"},
        )

        mock_response = MagicMock()
        mock_response.llm_output = {
            "token_usage": {
                "prompt_tokens": 0,
                "completion_tokens": 300,
                "total_tokens": 300,
            }
        }
        mock_response.generations = []

        callback.on_llm_end(mock_response)

        assert callback.counts["total_completion_tokens"] == 300

    def test_total_tokens_accumulation(self):
        """Test tokens accumulate across multiple calls"""
        callback = TokenCountingCallback()

        for i in range(3):
            callback.on_llm_start(
                {"_type": "ChatOpenAI"},
                ["test"],
                invocation_params={"model": "gpt-4"},
            )

            mock_response = MagicMock()
            mock_response.llm_output = {
                "token_usage": {
                    "prompt_tokens": 100,
                    "completion_tokens": 50,
                    "total_tokens": 150,
                }
            }
            mock_response.generations = []

            callback.on_llm_end(mock_response)

        assert callback.counts["total_tokens"] == 450
        assert callback.counts["by_model"]["gpt-4"]["calls"] == 3

    def test_context_overflow_detection(self):
        """Test context overflow detection at 95% threshold"""
        callback = TokenCountingCallback(
            research_context={"context_limit": 4096}
        )

        callback.on_llm_start(
            {"_type": "ChatOllama", "kwargs": {"model": "llama2"}}, ["test"]
        )

        # Simulate response with prompt_eval_count at exactly 95% of limit (3891.2 -> 3892)
        mock_generation = MagicMock()
        mock_generation.message.response_metadata = {
            "prompt_eval_count": 3900,  # ~95.2% of 4096 - above threshold
            "eval_count": 100,
        }
        # Set usage_metadata to None to force response_metadata path
        mock_generation.message.usage_metadata = None

        mock_response = MagicMock()
        mock_response.llm_output = {}
        mock_response.generations = [[mock_generation]]

        callback.on_llm_end(mock_response)

        assert callback.context_truncated is True

    def test_truncation_ratio_calculation(self):
        """Test truncation ratio is calculated correctly"""
        callback = TokenCountingCallback(
            research_context={"context_limit": 4096}
        )

        callback.on_llm_start(
            {"_type": "ChatOllama", "kwargs": {"model": "llama2"}},
            ["x" * 20000],
        )

        # Set original estimate AFTER on_llm_start so it doesn't get overwritten
        callback.original_prompt_estimate = 5000  # More than limit

        mock_generation = MagicMock()
        mock_generation.message.response_metadata = {
            "prompt_eval_count": 4000,  # ~97.6% of 4096 - above 95% threshold
            "eval_count": 50,
        }
        mock_generation.message.usage_metadata = None

        mock_response = MagicMock()
        mock_response.llm_output = {}
        mock_response.generations = [[mock_generation]]

        callback.on_llm_end(mock_response)

        # Truncation detected because we're at 97.6% of context limit
        assert callback.context_truncated is True

    def test_token_count_by_model(self):
        """Test tokens are tracked separately by model"""
        callback = TokenCountingCallback()

        # First call with gpt-4
        callback.on_llm_start(
            {"_type": "ChatOpenAI"},
            ["test"],
            invocation_params={"model": "gpt-4"},
        )

        mock_response = MagicMock()
        mock_response.llm_output = {
            "token_usage": {
                "prompt_tokens": 100,
                "completion_tokens": 50,
                "total_tokens": 150,
            }
        }
        mock_response.generations = []
        callback.on_llm_end(mock_response)

        # Second call with gpt-3.5
        callback.on_llm_start(
            {"_type": "ChatOpenAI"},
            ["test"],
            invocation_params={"model": "gpt-3.5-turbo"},
        )

        mock_response = MagicMock()
        mock_response.llm_output = {
            "token_usage": {
                "prompt_tokens": 200,
                "completion_tokens": 100,
                "total_tokens": 300,
            }
        }
        mock_response.generations = []
        callback.on_llm_end(mock_response)

        assert "gpt-4" in callback.counts["by_model"]
        assert "gpt-3.5-turbo" in callback.counts["by_model"]
        assert callback.counts["by_model"]["gpt-4"]["total_tokens"] == 150
        assert (
            callback.counts["by_model"]["gpt-3.5-turbo"]["total_tokens"] == 300
        )

    def test_token_count_by_provider(self):
        """Test provider is tracked for each model"""
        callback = TokenCountingCallback()

        callback.on_llm_start({"_type": "ChatOpenAI"}, ["test"])

        assert callback.current_provider == "openai"

    def test_ollama_metrics_parsing(self):
        """Test Ollama-specific metrics are parsed"""
        callback = TokenCountingCallback()

        callback.on_llm_start(
            {"_type": "ChatOllama", "kwargs": {"model": "llama2"}}, ["test"]
        )

        mock_generation = MagicMock()
        mock_generation.message.response_metadata = {
            "prompt_eval_count": 100,
            "eval_count": 50,
            "total_duration": 1000000000,
            "load_duration": 100000000,
            "prompt_eval_duration": 200000000,
            "eval_duration": 300000000,
        }
        mock_generation.message.usage_metadata = None

        mock_response = MagicMock()
        mock_response.llm_output = {}
        mock_response.generations = [[mock_generation]]

        callback.on_llm_end(mock_response)

        assert callback.ollama_metrics.get("prompt_eval_count") == 100
        assert callback.ollama_metrics.get("eval_count") == 50
        assert callback.ollama_metrics.get("total_duration") == 1000000000

    def test_anthropic_metrics_parsing(self):
        """Test Anthropic provider detection"""
        callback = TokenCountingCallback()

        callback.on_llm_start({"_type": "ChatAnthropic"}, ["test"])

        assert callback.current_provider == "anthropic"

    def test_openai_metrics_parsing(self):
        """Test OpenAI token usage parsing"""
        callback = TokenCountingCallback()

        callback.on_llm_start(
            {"_type": "ChatOpenAI"},
            ["test"],
            invocation_params={"model": "gpt-4"},
        )

        mock_response = MagicMock()
        mock_response.llm_output = {
            "usage": {
                "prompt_tokens": 100,
                "completion_tokens": 200,
                "total_tokens": 300,
            }
        }
        mock_response.generations = []

        callback.on_llm_end(mock_response)

        assert callback.counts["total_tokens"] == 300

    def test_unknown_provider_fallback(self):
        """Test unknown provider defaults correctly"""
        callback = TokenCountingCallback()

        callback.on_llm_start({"_type": "SomeUnknownLLM"}, ["test"])

        assert callback.current_provider == "unknown"

    def test_streaming_token_counting(self):
        """Test tokens are counted correctly in streaming mode"""
        callback = TokenCountingCallback()

        callback.on_llm_start(
            {"_type": "ChatOpenAI"},
            ["test"],
            invocation_params={"model": "gpt-4"},
        )

        # Streaming responses may have different token reporting
        mock_response = MagicMock()
        mock_response.llm_output = {
            "token_usage": {
                "prompt_tokens": 100,
                "completion_tokens": 50,
            }
        }
        mock_response.generations = []

        callback.on_llm_end(mock_response)

        # Should still extract tokens
        assert callback.counts["total_prompt_tokens"] == 100

    def test_batch_token_counting(self):
        """Test tokens counted for batch operations"""
        callback = TokenCountingCallback()

        callback.on_llm_start(
            {"_type": "ChatOpenAI"},
            ["test1", "test2", "test3"],
            invocation_params={"model": "gpt-4"},
        )

        mock_response = MagicMock()
        mock_response.llm_output = {
            "token_usage": {
                "prompt_tokens": 300,
                "completion_tokens": 150,
                "total_tokens": 450,
            }
        }
        mock_response.generations = []

        callback.on_llm_end(mock_response)

        assert callback.counts["total_tokens"] == 450


class TestCallStackTracking:
    """Tests for call stack tracking functionality"""

    def test_call_stack_pop(self):
        """Test call stack info is available after start"""
        callback = TokenCountingCallback()

        callback.on_llm_start({"_type": "ChatOpenAI"}, ["test"])

        # call_stack might be None if called from test context
        assert callback.call_stack is None or isinstance(
            callback.call_stack, str
        )

    def test_call_duration_tracking(self):
        """Test call duration is tracked"""
        callback = TokenCountingCallback()

        callback.on_llm_start({"_type": "ChatOpenAI"}, ["test"])
        time.sleep(0.01)  # 10ms

        mock_response = MagicMock()
        mock_response.llm_output = {}
        mock_response.generations = []

        callback.on_llm_end(mock_response)

        assert callback.response_time_ms is not None
        assert callback.response_time_ms >= 10

    def test_call_metadata_storage(self):
        """Test call metadata is stored"""
        callback = TokenCountingCallback(research_id="test-123")

        callback.on_llm_start({"_type": "ChatOpenAI"}, ["test"])

        assert callback.research_id == "test-123"

    def test_max_depth_tracking(self):
        """Test maximum stack depth is limited"""
        callback = TokenCountingCallback()

        callback.on_llm_start({"_type": "ChatOpenAI"}, ["test"])

        # Call stack should be limited to 5 frames
        if callback.call_stack:
            frames = callback.call_stack.split(" -> ")
            assert len(frames) <= 5

    def test_concurrent_call_tracking(self):
        """Test concurrent calls are tracked separately"""
        callback1 = TokenCountingCallback(research_id="test-1")
        callback2 = TokenCountingCallback(research_id="test-2")

        callback1.on_llm_start({"_type": "ChatOpenAI"}, ["test1"])
        callback2.on_llm_start({"_type": "ChatOpenAI"}, ["test2"])

        assert callback1.research_id != callback2.research_id

    def test_call_stack_serialization(self):
        """Test call stack can be serialized"""
        callback = TokenCountingCallback()

        callback.on_llm_start({"_type": "ChatOpenAI"}, ["test"])

        # Should be None or string (serializable)
        assert callback.call_stack is None or isinstance(
            callback.call_stack, str
        )


class TestResearchContextIntegration:
    """Tests for research context integration"""

    def test_research_context_mutation(self):
        """Test research context is properly stored"""
        context = {
            "query": "test query",
            "mode": "detailed",
            "iteration": 1,
        }
        callback = TokenCountingCallback(research_context=context)

        assert callback.research_context["query"] == "test query"
        assert callback.research_context["mode"] == "detailed"

    def test_context_token_budget(self):
        """Test context token budget is respected"""
        context = {"context_limit": 4096}
        callback = TokenCountingCallback(research_context=context)

        callback.on_llm_start({"_type": "ChatOllama"}, ["test"])

        assert callback.context_limit == 4096

    def test_context_overflow_handling(self):
        """Test context overflow is properly handled"""
        context = {"context_limit": 100}
        callback = TokenCountingCallback(research_context=context)

        callback.on_llm_start(
            {"_type": "ChatOllama", "kwargs": {"model": "llama2"}}, ["test"]
        )

        mock_generation = MagicMock()
        mock_generation.message.response_metadata = {
            "prompt_eval_count": 98,  # 98% of limit (above 95% threshold)
            "eval_count": 10,
        }
        mock_generation.message.usage_metadata = None

        mock_response = MagicMock()
        mock_response.llm_output = {}
        mock_response.generations = [[mock_generation]]

        callback.on_llm_end(mock_response)

        assert callback.context_truncated is True

    def test_context_truncation_strategy(self):
        """Test context truncation is detected"""
        callback = TokenCountingCallback(
            research_context={"context_limit": 1000}
        )

        callback.on_llm_start(
            {"_type": "ChatOllama", "kwargs": {"model": "llama2"}}, ["x" * 8000]
        )

        # Set original estimate AFTER on_llm_start
        callback.original_prompt_estimate = 2000

        mock_generation = MagicMock()
        mock_generation.message.response_metadata = {
            "prompt_eval_count": 960,  # 96% of limit - above 95% threshold
            "eval_count": 50,
        }
        mock_generation.message.usage_metadata = None

        mock_response = MagicMock()
        mock_response.llm_output = {}
        mock_response.generations = [[mock_generation]]

        callback.on_llm_end(mock_response)

        assert callback.context_truncated is True

    def test_context_priority_ordering(self):
        """Test research context priority is maintained"""
        context = {
            "priority": "high",
            "context_limit": 4096,
        }
        callback = TokenCountingCallback(research_context=context)

        assert callback.research_context.get("priority") == "high"

    def test_context_summary_generation(self):
        """Test get_counts returns summary"""
        callback = TokenCountingCallback()

        callback.on_llm_start(
            {"_type": "ChatOpenAI"},
            ["test"],
            invocation_params={"model": "gpt-4"},
        )

        mock_response = MagicMock()
        mock_response.llm_output = {
            "token_usage": {
                "prompt_tokens": 100,
                "completion_tokens": 50,
                "total_tokens": 150,
            }
        }
        mock_response.generations = []

        callback.on_llm_end(mock_response)

        counts = callback.counts

        assert "total_tokens" in counts
        assert "by_model" in counts

    def test_phase_token_allocation(self):
        """Test tokens tracked per research phase"""
        context = {"phase": "analysis"}
        callback = TokenCountingCallback(research_context=context)

        assert callback.research_context.get("phase") == "analysis"

    def test_synthesis_token_budget(self):
        """Test synthesis phase has token budget"""
        context = {"phase": "synthesis", "context_limit": 8192}
        callback = TokenCountingCallback(research_context=context)

        callback.on_llm_start({"_type": "ChatOpenAI"}, ["test"])

        assert callback.context_limit == 8192

    def test_analysis_token_budget(self):
        """Test analysis phase respects budget"""
        context = {"phase": "analysis", "context_limit": 4096}
        callback = TokenCountingCallback(research_context=context)

        callback.on_llm_start({"_type": "ChatOpenAI"}, ["test"])

        assert callback.context_limit == 4096

    def test_refinement_token_budget(self):
        """Test refinement phase has budget"""
        context = {"phase": "refinement", "context_limit": 2048}
        callback = TokenCountingCallback(research_context=context)

        callback.on_llm_start({"_type": "ChatOpenAI"}, ["test"])

        assert callback.context_limit == 2048

    def test_token_usage_reporting(self):
        """Test token usage can be reported"""
        callback = TokenCountingCallback()

        callback.on_llm_start(
            {"_type": "ChatOpenAI"},
            ["test"],
            invocation_params={"model": "gpt-4"},
        )

        mock_response = MagicMock()
        mock_response.llm_output = {
            "token_usage": {
                "prompt_tokens": 100,
                "completion_tokens": 50,
                "total_tokens": 150,
            }
        }
        mock_response.generations = []

        callback.on_llm_end(mock_response)

        counts = callback.counts
        assert counts["total_tokens"] == 150

    def test_cost_estimation_integration(self):
        """Test callback provides data for cost estimation"""
        callback = TokenCountingCallback()

        callback.on_llm_start(
            {"_type": "ChatOpenAI"},
            ["test"],
            invocation_params={"model": "gpt-4"},
        )

        mock_response = MagicMock()
        mock_response.llm_output = {
            "token_usage": {
                "prompt_tokens": 1000,
                "completion_tokens": 500,
                "total_tokens": 1500,
            }
        }
        mock_response.generations = []

        callback.on_llm_end(mock_response)

        # Data available for cost calculation
        assert callback.counts["total_prompt_tokens"] == 1000
        assert callback.counts["total_completion_tokens"] == 500
        assert callback.counts["by_model"]["gpt-4"]["provider"] == "openai"

    def test_model_pricing_lookup(self):
        """Test model name is available for pricing lookup"""
        callback = TokenCountingCallback()

        callback.preset_model = "gpt-4-turbo"
        callback.preset_provider = "openai"

        callback.on_llm_start({"_type": "ChatOpenAI"}, ["test"])

        assert callback.current_model == "gpt-4-turbo"
        assert callback.current_provider == "openai"

    def test_budget_exceeded_callback(self):
        """Test context overflow is flagged"""
        callback = TokenCountingCallback(
            research_context={"context_limit": 100}
        )

        callback.on_llm_start(
            {"_type": "ChatOllama", "kwargs": {"model": "llama2"}}, ["test"]
        )

        mock_generation = MagicMock()
        mock_generation.message.response_metadata = {
            "prompt_eval_count": 96,  # 96% of limit - above 95% threshold
            "eval_count": 10,
        }
        mock_generation.message.usage_metadata = None

        mock_response = MagicMock()
        mock_response.llm_output = {}
        mock_response.generations = [[mock_generation]]

        callback.on_llm_end(mock_response)

        assert callback.context_truncated is True

    def test_token_warning_thresholds(self):
        """Test warning at various threshold levels"""
        # Test at 70% - should not trigger
        callback_70 = TokenCountingCallback(
            research_context={"context_limit": 1000}
        )
        callback_70.on_llm_start(
            {"_type": "ChatOllama", "kwargs": {"model": "llama2"}}, ["test"]
        )

        mock_gen_70 = MagicMock()
        mock_gen_70.message.response_metadata = {
            "prompt_eval_count": 700,
            "eval_count": 10,
        }
        mock_gen_70.message.usage_metadata = None

        mock_resp_70 = MagicMock()
        mock_resp_70.llm_output = {}
        mock_resp_70.generations = [[mock_gen_70]]

        callback_70.on_llm_end(mock_resp_70)

        # 70% should not trigger (threshold is 80%)
        assert callback_70.context_truncated is False

        # Test at 96% - should trigger
        callback_96 = TokenCountingCallback(
            research_context={"context_limit": 1000}
        )
        callback_96.on_llm_start(
            {"_type": "ChatOllama", "kwargs": {"model": "llama2"}}, ["test"]
        )

        mock_gen_96 = MagicMock()
        mock_gen_96.message.response_metadata = {
            "prompt_eval_count": 960,
            "eval_count": 10,
        }
        mock_gen_96.message.usage_metadata = None

        mock_resp_96 = MagicMock()
        mock_resp_96.llm_output = {}
        mock_resp_96.generations = [[mock_gen_96]]

        callback_96.on_llm_end(mock_resp_96)

        assert callback_96.context_truncated is True


class TestErrorHandling:
    """Tests for error handling in token counter"""

    def test_on_llm_error_handling(self):
        """Test error handling callback"""
        callback = TokenCountingCallback()

        callback.on_llm_start({"_type": "ChatOpenAI"}, ["test"])

        # Simulate error
        callback.on_llm_error(Exception("Test error"))

        assert callback.success_status == "error"
        assert callback.error_type == "Exception"

    def test_missing_token_usage(self):
        """Test handling of missing token usage data"""
        callback = TokenCountingCallback()

        callback.on_llm_start({"_type": "ChatOpenAI"}, ["test"])

        mock_response = MagicMock()
        mock_response.llm_output = None
        mock_response.generations = []

        # Should not crash
        callback.on_llm_end(mock_response)

        # No tokens counted
        assert callback.counts["total_tokens"] == 0


class TestModelDetection:
    """Tests for model detection logic"""

    def test_model_from_invocation_params(self):
        """Test model extracted from invocation_params"""
        callback = TokenCountingCallback()

        callback.on_llm_start(
            {"_type": "ChatOpenAI"},
            ["test"],
            invocation_params={"model": "gpt-4-turbo"},
        )

        assert callback.current_model == "gpt-4-turbo"

    def test_model_from_serialized_kwargs(self):
        """Test model extracted from serialized kwargs"""
        callback = TokenCountingCallback()

        callback.on_llm_start(
            {"_type": "ChatOpenAI", "kwargs": {"model": "gpt-4o"}}, ["test"]
        )

        assert callback.current_model == "gpt-4o"

    def test_model_from_preset(self):
        """Test preset model is used"""
        callback = TokenCountingCallback()
        callback.preset_model = "custom-model"
        callback.preset_provider = "custom"

        callback.on_llm_start({"_type": "SomeLLM"}, ["test"])

        assert callback.current_model == "custom-model"
        assert callback.current_provider == "custom"

    def test_ollama_model_extraction(self):
        """Test Ollama model name extraction"""
        callback = TokenCountingCallback()

        callback.on_llm_start(
            {"_type": "ChatOllama", "kwargs": {"model": "llama2:7b"}}, ["test"]
        )

        assert callback.current_model == "llama2:7b"
        assert callback.current_provider == "ollama"

    def test_fallback_to_type_name(self):
        """Test fallback to type name when model not found"""
        callback = TokenCountingCallback()

        callback.on_llm_start({"_type": "UnknownLLMType"}, ["test"])

        # Should use the type as model name
        assert (
            "UnknownLLMType" in callback.current_model
            or callback.current_model == "unknown"
        )
