"""Test context overflow detection for LLM calls."""

import os
import uuid
from unittest.mock import Mock, patch

import pytest
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

from src.local_deep_research.database.models import Base, TokenUsage
from src.local_deep_research.metrics.token_counter import TokenCountingCallback


class TestContextOverflowDetection:
    """Test suite for context overflow detection."""

    @pytest.fixture
    def db_session(self):
        """Create an in-memory database session for testing."""
        engine = create_engine("sqlite:///:memory:")
        Base.metadata.create_all(engine)
        Session = sessionmaker(bind=engine)
        session = Session()
        yield session
        session.close()

    @pytest.fixture
    def token_callback(self):
        """Create a token counting callback for testing."""
        research_id = str(uuid.uuid4())
        research_context = {
            "research_query": "Test query",
            "research_mode": "test",
            "context_limit": 2048,  # Set a specific context limit
            "username": "test_user",
            "user_password": "test_pass",
        }
        return TokenCountingCallback(research_id, research_context)

    def test_context_overflow_detection_no_overflow(self, token_callback):
        """Test that no overflow is detected for small prompts."""
        # Simulate LLM start with small prompt
        prompts = ["What is 2+2?"]
        token_callback.on_llm_start({}, prompts)

        # Create mock response with Ollama-style metadata
        mock_response = Mock()
        mock_response.llm_output = None  # Explicitly set to None
        mock_response.generations = [[Mock()]]
        mock_response.generations[0][0].message = Mock()
        mock_response.generations[0][
            0
        ].message.usage_metadata = None  # No usage_metadata
        mock_response.generations[0][0].message.response_metadata = {
            "prompt_eval_count": 5,  # Small token count
            "eval_count": 10,
            "total_duration": 1000000000,  # 1 second in nanoseconds
        }

        # Process response
        token_callback.on_llm_end(mock_response)

        # Verify no overflow detected
        assert token_callback.context_truncated is False
        assert token_callback.tokens_truncated == 0
        assert token_callback.truncation_ratio == 0.0
        assert token_callback.ollama_metrics.get("prompt_eval_count") == 5

    def test_context_overflow_detection_with_overflow(self, token_callback):
        """Test that overflow is detected when prompt approaches context limit."""
        # Simulate LLM start with large prompt
        large_text = "The quick brown fox jumps over the lazy dog. " * 500
        prompts = [large_text]
        token_callback.on_llm_start({}, prompts)

        # Create mock response indicating near-limit token usage
        mock_response = Mock()
        mock_response.llm_output = None
        mock_response.generations = [[Mock()]]
        mock_response.generations[0][0].message = Mock()
        mock_response.generations[0][0].message.usage_metadata = None
        mock_response.generations[0][0].message.response_metadata = {
            "prompt_eval_count": 1950,  # Near the 2048 limit (95%)
            "eval_count": 50,
            "total_duration": 5000000000,  # 5 seconds
            "prompt_eval_duration": 4000000000,
            "eval_duration": 1000000000,
        }

        # Process response
        token_callback.on_llm_end(mock_response)

        # Verify overflow detected
        assert token_callback.context_truncated is True
        assert token_callback.tokens_truncated > 0  # Should estimate truncation
        assert token_callback.truncation_ratio > 0
        assert token_callback.ollama_metrics["prompt_eval_count"] == 1950

    def test_ollama_raw_metrics_capture(self, token_callback):
        """Test that raw Ollama metrics are properly captured."""
        prompts = ["Test prompt"]
        token_callback.on_llm_start({}, prompts)

        # Create mock response with full Ollama metrics
        mock_response = Mock()
        mock_response.llm_output = None
        mock_response.generations = [[Mock()]]
        mock_response.generations[0][0].message = Mock()
        mock_response.generations[0][0].message.usage_metadata = None
        mock_response.generations[0][0].message.response_metadata = {
            "prompt_eval_count": 100,
            "eval_count": 200,
            "total_duration": 3000000000,
            "load_duration": 500000000,
            "prompt_eval_duration": 2000000000,
            "eval_duration": 500000000,
        }

        # Process response
        token_callback.on_llm_end(mock_response)

        # Verify all metrics captured
        assert token_callback.ollama_metrics["prompt_eval_count"] == 100
        assert token_callback.ollama_metrics["eval_count"] == 200
        assert token_callback.ollama_metrics["total_duration"] == 3000000000
        assert token_callback.ollama_metrics["load_duration"] == 500000000
        assert (
            token_callback.ollama_metrics["prompt_eval_duration"] == 2000000000
        )
        assert token_callback.ollama_metrics["eval_duration"] == 500000000

    def test_context_limit_from_research_context(self):
        """Test that context limit is properly read from research context."""
        # Create callback with context limit
        callback = TokenCountingCallback("test-id", {"context_limit": 2048})
        # Context limit is set on llm_start
        callback.on_llm_start({}, ["test"])
        assert callback.context_limit == 2048

        # Test with different limit
        callback_4k = TokenCountingCallback("test-id", {"context_limit": 4096})
        callback_4k.on_llm_start({}, ["test"])
        assert callback_4k.context_limit == 4096

        # Test with no limit
        callback_no_limit = TokenCountingCallback("test-id", {})
        callback_no_limit.on_llm_start({}, ["test"])
        assert callback_no_limit.context_limit is None

    def test_prompt_size_estimation(self, token_callback):
        """Test that prompt size is estimated correctly."""
        # Test with single prompt
        prompts = ["This is a test prompt with approximately 10 words."]
        token_callback.on_llm_start({}, prompts)

        # Rough estimate: ~4 chars per token
        expected_tokens = len(prompts[0]) // 4
        assert (
            abs(token_callback.original_prompt_estimate - expected_tokens) < 5
        )

        # Test with multiple prompts
        token_callback.original_prompt_estimate = 0
        prompts = ["First prompt.", "Second prompt.", "Third prompt."]
        total_chars = sum(len(p) for p in prompts)
        token_callback.on_llm_start({}, prompts)

        expected_tokens = total_chars // 4
        assert (
            abs(token_callback.original_prompt_estimate - expected_tokens) < 5
        )

    @patch("src.local_deep_research.metrics.token_counter.logger")
    def test_overflow_warning_logged(self, mock_logger, token_callback):
        """Test that overflow detection logs a warning."""
        # Create large prompt
        large_text = "word " * 10000
        prompts = [large_text]
        token_callback.on_llm_start({}, prompts)

        # Mock response at context limit
        mock_response = Mock()
        mock_response.llm_output = None
        mock_response.generations = [[Mock()]]
        mock_response.generations[0][0].message = Mock()
        mock_response.generations[0][0].message.usage_metadata = None
        mock_response.generations[0][0].message.response_metadata = {
            "prompt_eval_count": 2000,  # At 95% of 2048 limit
            "eval_count": 10,
        }

        token_callback.on_llm_end(mock_response)

        # Verify warning was logged
        mock_logger.warning.assert_called()
        warning_call = mock_logger.warning.call_args[0][0]
        assert "Context overflow detected" in warning_call
        assert "2000/2048" in warning_call

    def test_token_usage_save_with_overflow_fields(
        self, db_session, monkeypatch
    ):
        """Test that overflow fields are saved to database."""
        # Skip this test as it requires Flask context
        pytest.skip("Requires Flask application context")

        # Mock get_user_db_session to return our test session
        def mock_get_session(username):
            class SessionContext:
                def __enter__(self):
                    return db_session

                def __exit__(self, *args):
                    pass

            return SessionContext()

        monkeypatch.setattr(
            "src.local_deep_research.metrics.token_counter.get_user_db_session",
            mock_get_session,
        )

        # Create callback and simulate overflow
        research_id = str(uuid.uuid4())
        callback = TokenCountingCallback(research_id, {"context_limit": 2048})

        # Simulate overflow scenario
        callback.current_model = "test-model"
        callback.current_provider = "ollama"
        callback.context_truncated = True
        callback.tokens_truncated = 500
        callback.truncation_ratio = 0.2
        callback.ollama_metrics = {
            "prompt_eval_count": 2000,
            "eval_count": 100,
            "total_duration": 5000000000,
        }

        # Save to database
        callback._save_to_db(2000, 100)

        # Verify saved data
        saved = (
            db_session.query(TokenUsage)
            .filter_by(research_id=research_id)
            .first()
        )
        assert saved is not None
        assert saved.context_limit == 2048
        assert saved.context_truncated == 1
        assert saved.tokens_truncated == 500
        assert saved.truncation_ratio == 0.2
        assert saved.ollama_prompt_eval_count == 2000
        assert saved.ollama_eval_count == 100
        assert saved.ollama_total_duration == 5000000000


@pytest.mark.skipif(
    os.environ.get("SKIP_OLLAMA_TESTS", "true").lower() == "true",
    reason="Ollama integration tests skipped",
)
class TestContextOverflowIntegration:
    """Integration tests with actual Ollama (when available)."""

    @pytest.mark.slow
    def test_ollama_context_overflow_real(self):
        """Test with real Ollama instance if available."""
        from langchain_ollama import ChatOllama
        from src.local_deep_research.config.llm_config import (
            is_ollama_available,
        )

        if not is_ollama_available():
            pytest.skip("Ollama not available")

        # Create LLM with small context window
        llm = ChatOllama(
            model="llama3.2:latest",
            num_ctx=512,  # Very small context for testing
            temperature=0.1,
        )

        # Create callback
        research_id = str(uuid.uuid4())
        callback = TokenCountingCallback(research_id, {"context_limit": 512})

        # Create prompt that will likely overflow
        large_prompt = "Please analyze this text: " + ("word " * 200)

        # Run with callback
        try:
            _ = llm.invoke(large_prompt, config={"callbacks": [callback]})

            # Check if overflow was detected
            if callback.ollama_metrics.get("prompt_eval_count"):
                prompt_tokens = callback.ollama_metrics["prompt_eval_count"]
                if prompt_tokens >= 512 * 0.95:
                    assert callback.context_truncated is True
                    print(f"✅ Overflow detected: {prompt_tokens}/512 tokens")
                else:
                    print(f"ℹ️ No overflow: {prompt_tokens}/512 tokens")
        except Exception as e:
            pytest.skip(f"Ollama test failed: {e}")
