"""
Comprehensive tests for token_counter.py metrics aggregation.

Tests cover:
- _get_metrics_from_encrypted_db with various time periods
- _get_metrics_from_thread_db unit tests
- Context overflow detection edge cases
- Ollama-specific metrics (prompt_eval_count, eval_count handling)
- Call stack tracking integration
- Research context filtering
- Rate limiting metrics aggregation
- Model cost calculations
- Metrics merging
"""

import time
from unittest.mock import MagicMock, Mock, patch


class TestGetMetricsFromEncryptedDb:
    """Tests for _get_metrics_from_encrypted_db functionality."""

    def test_no_username_returns_empty_metrics(self):
        """Test empty metrics returned when no username in session."""
        from local_deep_research.metrics.token_counter import TokenCounter

        counter = TokenCounter()

        with patch("flask.session", {}):
            result = counter._get_metrics_from_encrypted_db("30d", "all")

            assert result["total_tokens"] == 0
            assert result["total_researches"] == 0
            assert result["by_model"] == []

    def test_time_filter_7d(self):
        """Test 7-day time filter is applied correctly."""
        from local_deep_research.metrics.token_counter import TokenCounter

        counter = TokenCounter()

        mock_session = MagicMock()
        mock_query = MagicMock()
        mock_session.query.return_value = mock_query
        mock_query.filter.return_value = mock_query
        mock_query.with_entities.return_value = mock_query
        mock_query.scalar.return_value = 1000
        mock_query.group_by.return_value = mock_query
        mock_query.order_by.return_value = mock_query
        mock_query.limit.return_value = mock_query
        mock_query.all.return_value = []
        mock_query.first.return_value = MagicMock(
            total_input_tokens=500,
            total_output_tokens=500,
            avg_input_tokens=50,
            avg_output_tokens=50,
            avg_total_tokens=100,
        )
        mock_query.count.return_value = 0

        with patch("flask.session", {"username": "testuser"}):
            with patch(
                "local_deep_research.database.session_context.get_user_db_session"
            ) as mock_get_session:
                mock_get_session.return_value.__enter__ = Mock(
                    return_value=mock_session
                )
                mock_get_session.return_value.__exit__ = Mock(
                    return_value=False
                )

                counter._get_metrics_from_encrypted_db("7d", "all")

                # Verify filter was called (time filter applied)
                assert mock_query.filter.called

    def test_time_filter_30d(self):
        """Test 30-day time filter is applied correctly."""
        from local_deep_research.metrics.token_counter import TokenCounter

        counter = TokenCounter()

        mock_session = MagicMock()
        mock_query = MagicMock()
        mock_session.query.return_value = mock_query
        mock_query.filter.return_value = mock_query
        mock_query.with_entities.return_value = mock_query
        mock_query.scalar.return_value = 5000
        mock_query.group_by.return_value = mock_query
        mock_query.order_by.return_value = mock_query
        mock_query.limit.return_value = mock_query
        mock_query.all.return_value = []
        mock_query.first.return_value = MagicMock(
            total_input_tokens=2500,
            total_output_tokens=2500,
            avg_input_tokens=100,
            avg_output_tokens=100,
            avg_total_tokens=200,
        )
        mock_query.count.return_value = 0

        with patch("flask.session", {"username": "testuser"}):
            with patch(
                "local_deep_research.database.session_context.get_user_db_session"
            ) as mock_get_session:
                mock_get_session.return_value.__enter__ = Mock(
                    return_value=mock_session
                )
                mock_get_session.return_value.__exit__ = Mock(
                    return_value=False
                )

                counter._get_metrics_from_encrypted_db("30d", "all")

                assert mock_query.filter.called

    def test_time_filter_all(self):
        """Test 'all' time filter returns all data without time filtering."""
        from local_deep_research.metrics.token_counter import TokenCounter

        counter = TokenCounter()

        mock_session = MagicMock()
        mock_query = MagicMock()
        mock_session.query.return_value = mock_query
        mock_query.filter.return_value = mock_query
        mock_query.with_entities.return_value = mock_query
        mock_query.scalar.return_value = 10000
        mock_query.group_by.return_value = mock_query
        mock_query.order_by.return_value = mock_query
        mock_query.limit.return_value = mock_query
        mock_query.all.return_value = []
        mock_query.first.return_value = MagicMock(
            total_input_tokens=5000,
            total_output_tokens=5000,
            avg_input_tokens=200,
            avg_output_tokens=200,
            avg_total_tokens=400,
        )
        mock_query.count.return_value = 0

        # Use a context manager mock properly
        mock_context = MagicMock()
        mock_context.__enter__ = Mock(return_value=mock_session)
        mock_context.__exit__ = Mock(return_value=False)

        with patch("flask.session", {"username": "testuser"}):
            with patch(
                "local_deep_research.database.session_context.get_user_db_session",
                return_value=mock_context,
            ):
                result = counter._get_metrics_from_encrypted_db("all", "all")

                # Verify returns a metrics dict with expected structure
                assert "total_tokens" in result
                assert "total_researches" in result
                assert "by_model" in result

    def test_research_mode_filter_quick(self):
        """Test 'quick' research mode filter is applied."""
        from local_deep_research.metrics.token_counter import TokenCounter

        counter = TokenCounter()

        mock_session = MagicMock()
        mock_query = MagicMock()
        mock_session.query.return_value = mock_query
        mock_query.filter.return_value = mock_query
        mock_query.with_entities.return_value = mock_query
        mock_query.scalar.return_value = 500
        mock_query.group_by.return_value = mock_query
        mock_query.order_by.return_value = mock_query
        mock_query.limit.return_value = mock_query
        mock_query.all.return_value = []
        mock_query.first.return_value = MagicMock(
            total_input_tokens=250,
            total_output_tokens=250,
            avg_input_tokens=25,
            avg_output_tokens=25,
            avg_total_tokens=50,
        )
        mock_query.count.return_value = 0

        with patch("flask.session", {"username": "testuser"}):
            with patch(
                "local_deep_research.database.session_context.get_user_db_session"
            ) as mock_get_session:
                mock_get_session.return_value.__enter__ = Mock(
                    return_value=mock_session
                )
                mock_get_session.return_value.__exit__ = Mock(
                    return_value=False
                )

                counter._get_metrics_from_encrypted_db("30d", "quick")

                # Mode filter should be applied
                assert mock_query.filter.called

    def test_database_error_returns_empty_metrics(self):
        """Test database error returns empty metrics."""
        from local_deep_research.metrics.token_counter import TokenCounter

        counter = TokenCounter()

        with patch("flask.session", {"username": "testuser"}):
            with patch(
                "local_deep_research.database.session_context.get_user_db_session"
            ) as mock_get_session:
                mock_get_session.side_effect = Exception("Database error")

                result = counter._get_metrics_from_encrypted_db("30d", "all")

                assert result["total_tokens"] == 0
                assert result["total_researches"] == 0


class TestGetMetricsFromThreadDb:
    """Tests for _get_metrics_from_thread_db functionality."""

    def test_no_thread_metrics_db_returns_empty(self):
        """Test empty metrics when thread metrics DB not available."""
        from local_deep_research.metrics.token_counter import TokenCounter

        counter = TokenCounter()
        counter._thread_metrics_db = None

        result = counter._get_metrics_from_thread_db("30d", "all")

        assert result["total_tokens"] == 0
        assert result["total_researches"] == 0
        assert result["by_model"] == []

    def test_thread_db_returns_metrics(self):
        """Test thread DB returns proper metrics."""
        from local_deep_research.metrics.token_counter import TokenCounter

        counter = TokenCounter()

        mock_session = MagicMock()
        mock_query = MagicMock()
        mock_session.query.return_value = mock_query
        mock_query.filter.return_value = mock_query
        mock_query.with_entities.return_value = mock_query
        mock_query.scalar.return_value = 2000
        mock_query.group_by.return_value = mock_query
        mock_query.all.return_value = []
        mock_query.first.return_value = MagicMock(
            total_input_tokens=1000,
            total_output_tokens=1000,
        )

        mock_thread_db = MagicMock()
        mock_thread_db.get_session.return_value.__enter__ = Mock(
            return_value=mock_session
        )
        mock_thread_db.get_session.return_value.__exit__ = Mock(
            return_value=False
        )

        counter._thread_metrics_db = mock_thread_db

        result = counter._get_metrics_from_thread_db("30d", "all")

        assert result["total_tokens"] == 2000

    def test_thread_db_error_returns_empty(self):
        """Test thread DB error returns empty metrics."""
        from local_deep_research.metrics.token_counter import TokenCounter

        counter = TokenCounter()

        mock_thread_db = MagicMock()
        mock_thread_db.get_session.side_effect = Exception("Thread DB error")

        counter._thread_metrics_db = mock_thread_db

        result = counter._get_metrics_from_thread_db("30d", "all")

        assert result["total_tokens"] == 0


class TestMergeMetrics:
    """Tests for _merge_metrics functionality."""

    def test_merge_empty_metrics(self):
        """Test merging empty metrics."""
        from local_deep_research.metrics.token_counter import TokenCounter

        counter = TokenCounter()

        encrypted = {
            "total_tokens": 0,
            "total_researches": 0,
            "by_model": [],
            "recent_researches": [],
            "token_breakdown": {},
        }
        thread = {
            "total_tokens": 0,
            "total_researches": 0,
            "by_model": [],
            "recent_researches": [],
            "token_breakdown": {},
        }

        result = counter._merge_metrics(encrypted, thread)

        assert result["total_tokens"] == 0
        assert result["total_researches"] == 0

    def test_merge_adds_tokens(self):
        """Test merging adds tokens from both sources."""
        from local_deep_research.metrics.token_counter import TokenCounter

        counter = TokenCounter()

        encrypted = {
            "total_tokens": 1000,
            "total_researches": 5,
            "by_model": [],
            "recent_researches": [],
            "token_breakdown": {
                "total_input_tokens": 500,
                "total_output_tokens": 500,
            },
        }
        thread = {
            "total_tokens": 2000,
            "total_researches": 3,
            "by_model": [],
            "recent_researches": [],
            "token_breakdown": {
                "total_input_tokens": 1000,
                "total_output_tokens": 1000,
            },
        }

        result = counter._merge_metrics(encrypted, thread)

        assert result["total_tokens"] == 3000
        # Takes max of researches
        assert result["total_researches"] == 5

    def test_merge_model_usage(self):
        """Test merging model usage data."""
        from local_deep_research.metrics.token_counter import TokenCounter

        counter = TokenCounter()

        encrypted = {
            "total_tokens": 1000,
            "total_researches": 5,
            "by_model": [
                {
                    "model": "gpt-4",
                    "provider": "openai",
                    "tokens": 500,
                    "calls": 5,
                    "prompt_tokens": 250,
                    "completion_tokens": 250,
                }
            ],
            "recent_researches": [],
            "token_breakdown": {},
        }
        thread = {
            "total_tokens": 500,
            "total_researches": 3,
            "by_model": [
                {
                    "model": "gpt-4",
                    "provider": "openai",
                    "tokens": 300,
                    "calls": 3,
                    "prompt_tokens": 150,
                    "completion_tokens": 150,
                }
            ],
            "recent_researches": [],
            "token_breakdown": {},
        }

        result = counter._merge_metrics(encrypted, thread)

        # Should merge gpt-4 usage
        assert len(result["by_model"]) == 1
        assert result["by_model"][0]["model"] == "gpt-4"
        assert result["by_model"][0]["tokens"] == 800
        assert result["by_model"][0]["calls"] == 8

    def test_merge_different_models(self):
        """Test merging keeps different models separate."""
        from local_deep_research.metrics.token_counter import TokenCounter

        counter = TokenCounter()

        encrypted = {
            "total_tokens": 1000,
            "total_researches": 5,
            "by_model": [
                {
                    "model": "gpt-4",
                    "provider": "openai",
                    "tokens": 500,
                    "calls": 5,
                    "prompt_tokens": 250,
                    "completion_tokens": 250,
                }
            ],
            "recent_researches": [],
            "token_breakdown": {},
        }
        thread = {
            "total_tokens": 500,
            "total_researches": 3,
            "by_model": [
                {
                    "model": "claude-3",
                    "provider": "anthropic",
                    "tokens": 300,
                    "calls": 3,
                    "prompt_tokens": 150,
                    "completion_tokens": 150,
                }
            ],
            "recent_researches": [],
            "token_breakdown": {},
        }

        result = counter._merge_metrics(encrypted, thread)

        assert len(result["by_model"]) == 2

    def test_merge_token_breakdown(self):
        """Test merging token breakdown data."""
        from local_deep_research.metrics.token_counter import TokenCounter

        counter = TokenCounter()

        encrypted = {
            "total_tokens": 1000,
            "total_researches": 5,
            "by_model": [],
            "recent_researches": [],
            "token_breakdown": {
                "total_input_tokens": 500,
                "total_output_tokens": 500,
                "avg_input_tokens": 50,
                "avg_output_tokens": 50,
                "avg_total_tokens": 100,
            },
        }
        thread = {
            "total_tokens": 500,
            "total_researches": 3,
            "by_model": [],
            "recent_researches": [],
            "token_breakdown": {
                "total_input_tokens": 250,
                "total_output_tokens": 250,
            },
        }

        result = counter._merge_metrics(encrypted, thread)

        assert result["token_breakdown"]["total_input_tokens"] == 750
        assert result["token_breakdown"]["total_output_tokens"] == 750


class TestContextOverflowEdgeCases:
    """Tests for context overflow detection edge cases."""

    def test_context_overflow_at_exactly_95_percent(self):
        """Test context overflow at exactly 95% threshold."""
        from local_deep_research.metrics.token_counter import (
            TokenCountingCallback,
        )

        callback = TokenCountingCallback(
            research_context={"context_limit": 1000}
        )

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

        mock_generation = MagicMock()
        mock_generation.message.response_metadata = {
            "prompt_eval_count": 950,  # Exactly 95%
            "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_overflow_at_94_percent(self):
        """Test context overflow not triggered at 94%."""
        from local_deep_research.metrics.token_counter import (
            TokenCountingCallback,
        )

        callback = TokenCountingCallback(
            research_context={"context_limit": 1000}
        )

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

        mock_generation = MagicMock()
        mock_generation.message.response_metadata = {
            "prompt_eval_count": 940,  # 94% - below 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 False

    def test_context_overflow_no_limit_set(self):
        """Test context overflow with no limit set."""
        from local_deep_research.metrics.token_counter import (
            TokenCountingCallback,
        )

        callback = TokenCountingCallback(research_context={})

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

        mock_generation = MagicMock()
        mock_generation.message.response_metadata = {
            "prompt_eval_count": 10000,
            "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)

        # No truncation without context limit
        assert callback.context_truncated is False

    def test_context_overflow_zero_prompt_count(self):
        """Test context overflow handling with zero prompt count."""
        from local_deep_research.metrics.token_counter import (
            TokenCountingCallback,
        )

        callback = TokenCountingCallback(
            research_context={"context_limit": 1000}
        )

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

        mock_generation = MagicMock()
        mock_generation.message.response_metadata = {
            "prompt_eval_count": 0,
            "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 False


class TestOllamaSpecificMetrics:
    """Tests for Ollama-specific metrics handling."""

    def test_ollama_prompt_eval_count(self):
        """Test Ollama prompt_eval_count is captured."""
        from local_deep_research.metrics.token_counter import (
            TokenCountingCallback,
        )

        callback = TokenCountingCallback()

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

        mock_generation = MagicMock()
        mock_generation.message.response_metadata = {
            "prompt_eval_count": 150,
            "eval_count": 75,
            "total_duration": 2000000000,
            "load_duration": 100000000,
            "prompt_eval_duration": 500000000,
            "eval_duration": 1000000000,
        }
        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["prompt_eval_count"] == 150
        assert callback.ollama_metrics["eval_count"] == 75
        assert callback.ollama_metrics["total_duration"] == 2000000000

    def test_ollama_metrics_stored_for_db(self):
        """Test Ollama metrics are included in context overflow fields."""
        from local_deep_research.metrics.token_counter import (
            TokenCountingCallback,
        )

        callback = TokenCountingCallback()

        callback.ollama_metrics = {
            "prompt_eval_count": 100,
            "eval_count": 50,
            "total_duration": 1000000,
            "load_duration": 100000,
            "prompt_eval_duration": 300000,
            "eval_duration": 400000,
        }

        fields = callback._get_context_overflow_fields()

        assert fields["ollama_prompt_eval_count"] == 100
        assert fields["ollama_eval_count"] == 50
        assert fields["ollama_total_duration"] == 1000000

    def test_ollama_provider_detection(self):
        """Test Ollama provider is correctly detected."""
        from local_deep_research.metrics.token_counter import (
            TokenCountingCallback,
        )

        callback = TokenCountingCallback()

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

        assert callback.current_provider == "ollama"
        assert callback.current_model == "mistral"


class TestCallStackTracking:
    """Tests for call stack tracking integration."""

    def test_call_stack_captured(self):
        """Test call stack is captured on LLM start."""
        from local_deep_research.metrics.token_counter import (
            TokenCountingCallback,
        )

        callback = TokenCountingCallback()

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

        # Call stack attributes should exist
        assert hasattr(callback, "calling_file")
        assert hasattr(callback, "calling_function")
        assert hasattr(callback, "call_stack")

    def test_call_stack_limited_to_5_frames(self):
        """Test call stack is limited to 5 frames."""
        from local_deep_research.metrics.token_counter import (
            TokenCountingCallback,
        )

        callback = TokenCountingCallback()

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

        if callback.call_stack:
            frames = callback.call_stack.split(" -> ")
            assert len(frames) <= 5

    def test_calling_file_extraction(self):
        """Test calling file is extracted correctly."""
        from local_deep_research.metrics.token_counter import (
            TokenCountingCallback,
        )

        callback = TokenCountingCallback()

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

        # In test context, calling_file might be None
        # Just verify the attribute exists
        assert hasattr(callback, "calling_file")


class TestResearchContextFiltering:
    """Tests for research context filtering."""

    def test_research_phase_stored(self):
        """Test research phase is stored in context."""
        from local_deep_research.metrics.token_counter import (
            TokenCountingCallback,
        )

        context = {"research_phase": "analysis"}
        callback = TokenCountingCallback(research_context=context)

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

    def test_search_iteration_stored(self):
        """Test search iteration is stored in context."""
        from local_deep_research.metrics.token_counter import (
            TokenCountingCallback,
        )

        context = {"search_iteration": 3}
        callback = TokenCountingCallback(research_context=context)

        assert callback.research_context.get("search_iteration") == 3

    def test_research_query_stored(self):
        """Test research query is stored in context."""
        from local_deep_research.metrics.token_counter import (
            TokenCountingCallback,
        )

        context = {"research_query": "What is machine learning?"}
        callback = TokenCountingCallback(research_context=context)

        assert (
            callback.research_context.get("research_query")
            == "What is machine learning?"
        )

    def test_research_mode_stored(self):
        """Test research mode is stored in context."""
        from local_deep_research.metrics.token_counter import (
            TokenCountingCallback,
        )

        context = {"research_mode": "detailed"}
        callback = TokenCountingCallback(research_context=context)

        assert callback.research_context.get("research_mode") == "detailed"


class TestGetEmptyMetrics:
    """Tests for _get_empty_metrics functionality."""

    def test_empty_metrics_structure(self):
        """Test empty metrics has correct structure."""
        from local_deep_research.metrics.token_counter import TokenCounter

        counter = TokenCounter()

        result = counter._get_empty_metrics()

        assert "total_tokens" in result
        assert "total_researches" in result
        assert "by_model" in result
        assert "recent_researches" in result
        assert "token_breakdown" in result
        assert result["total_tokens"] == 0
        assert result["total_researches"] == 0


class TestGetOverallMetrics:
    """Tests for get_overall_metrics functionality."""

    def test_overall_metrics_merges_sources(self):
        """Test overall metrics merges encrypted and thread DB sources."""
        from local_deep_research.metrics.token_counter import TokenCounter

        counter = TokenCounter()

        encrypted_metrics = {
            "total_tokens": 1000,
            "total_researches": 5,
            "by_model": [],
            "recent_researches": [],
            "token_breakdown": {},
        }
        thread_metrics = {
            "total_tokens": 500,
            "total_researches": 3,
            "by_model": [],
            "recent_researches": [],
            "token_breakdown": {},
        }

        with patch.object(
            counter,
            "_get_metrics_from_encrypted_db",
            return_value=encrypted_metrics,
        ):
            with patch.object(
                counter,
                "_get_metrics_from_thread_db",
                return_value=thread_metrics,
            ):
                result = counter.get_overall_metrics("30d", "all")

                assert result["total_tokens"] == 1500


class TestCreateCallback:
    """Tests for create_callback functionality."""

    def test_create_callback_with_research_id(self):
        """Test callback creation with research ID."""
        from local_deep_research.metrics.token_counter import TokenCounter

        counter = TokenCounter()

        callback = counter.create_callback(research_id="test-123")

        assert callback.research_id == "test-123"

    def test_create_callback_with_research_context(self):
        """Test callback creation with research context."""
        from local_deep_research.metrics.token_counter import TokenCounter

        counter = TokenCounter()

        context = {"research_mode": "detailed", "context_limit": 4096}
        callback = counter.create_callback(research_context=context)

        assert callback.research_context["research_mode"] == "detailed"
        assert callback.research_context["context_limit"] == 4096

    def test_create_callback_default_values(self):
        """Test callback creation with default values."""
        from local_deep_research.metrics.token_counter import TokenCounter

        counter = TokenCounter()

        callback = counter.create_callback()

        assert callback.research_id is None
        assert callback.research_context == {}


class TestGetResearchMetrics:
    """Tests for get_research_metrics functionality."""

    def test_no_username_returns_empty(self):
        """Test empty metrics when no username in session."""
        from local_deep_research.metrics.token_counter import TokenCounter

        counter = TokenCounter()

        with patch("flask.session", {}):
            result = counter.get_research_metrics("test-research-id")

            assert result["total_tokens"] == 0
            assert result["total_calls"] == 0


class TestResponseTimeTracking:
    """Tests for response time tracking."""

    def test_response_time_calculated(self):
        """Test response time is calculated correctly."""
        from local_deep_research.metrics.token_counter import (
            TokenCountingCallback,
        )

        callback = TokenCountingCallback()

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

        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 >= 50

    def test_response_time_on_error(self):
        """Test response time is tracked even on error."""
        from local_deep_research.metrics.token_counter import (
            TokenCountingCallback,
        )

        callback = TokenCountingCallback()

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

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

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


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

    def test_error_status_tracked(self):
        """Test error status is tracked."""
        from local_deep_research.metrics.token_counter import (
            TokenCountingCallback,
        )

        callback = TokenCountingCallback()

        callback.on_llm_start({"_type": "ChatOpenAI"}, ["test"])
        callback.on_llm_error(ValueError("Invalid input"))

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

    def test_success_status_default(self):
        """Test success status is default."""
        from local_deep_research.metrics.token_counter import (
            TokenCountingCallback,
        )

        callback = TokenCountingCallback()

        assert callback.success_status == "success"
        assert callback.error_type is None


class TestPresetModelAndProvider:
    """Tests for preset model and provider functionality."""

    def test_preset_model_used(self):
        """Test preset model is used over detected model."""
        from local_deep_research.metrics.token_counter import (
            TokenCountingCallback,
        )

        callback = TokenCountingCallback()
        callback.preset_model = "custom-model"
        callback.preset_provider = "custom-provider"

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

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

    def test_detected_model_when_no_preset(self):
        """Test detected model when no preset is set."""
        from local_deep_research.metrics.token_counter import (
            TokenCountingCallback,
        )

        callback = TokenCountingCallback()

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

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


class TestTokenUsageExtraction:
    """Tests for token usage extraction from various response formats."""

    def test_extraction_from_llm_output_token_usage(self):
        """Test extraction from llm_output.token_usage."""
        from local_deep_research.metrics.token_counter import (
            TokenCountingCallback,
        )

        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)

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

    def test_extraction_from_llm_output_usage(self):
        """Test extraction from llm_output.usage."""
        from local_deep_research.metrics.token_counter import (
            TokenCountingCallback,
        )

        callback = TokenCountingCallback()

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

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

        callback.on_llm_end(mock_response)

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

    def test_extraction_from_generation_usage_metadata(self):
        """Test extraction from generation.message.usage_metadata."""
        from local_deep_research.metrics.token_counter import (
            TokenCountingCallback,
        )

        callback = TokenCountingCallback()

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

        mock_generation = MagicMock()
        mock_generation.message.usage_metadata = {
            "input_tokens": 75,
            "output_tokens": 25,
            "total_tokens": 100,
        }

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

        callback.on_llm_end(mock_response)

        assert callback.counts["total_prompt_tokens"] == 75
        assert callback.counts["total_completion_tokens"] == 25

    def test_no_token_usage_available(self):
        """Test handling when no token usage is available."""
        from local_deep_research.metrics.token_counter import (
            TokenCountingCallback,
        )

        callback = TokenCountingCallback()

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

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

        callback.on_llm_end(mock_response)

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


class TestGetCounts:
    """Tests for get_counts functionality."""

    def test_get_counts_returns_counts(self):
        """Test get_counts returns current counts."""
        from local_deep_research.metrics.token_counter import (
            TokenCountingCallback,
        )

        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.get_counts()

        assert counts["total_tokens"] == 150
        assert counts["total_prompt_tokens"] == 100
        assert counts["total_completion_tokens"] == 50
        assert "gpt-4" in counts["by_model"]


class TestSearchEngineContext:
    """Tests for search engine context tracking."""

    def test_search_engine_planned_stored(self):
        """Test search engines planned is stored."""
        from local_deep_research.metrics.token_counter import (
            TokenCountingCallback,
        )

        context = {"search_engines_planned": ["google", "bing", "duckduckgo"]}
        callback = TokenCountingCallback(research_context=context)

        assert callback.research_context.get("search_engines_planned") == [
            "google",
            "bing",
            "duckduckgo",
        ]

    def test_search_engine_selected_stored(self):
        """Test search engine selected is stored."""
        from local_deep_research.metrics.token_counter import (
            TokenCountingCallback,
        )

        context = {"search_engine_selected": "google"}
        callback = TokenCountingCallback(research_context=context)

        assert (
            callback.research_context.get("search_engine_selected") == "google"
        )
