"""Tests for ContextRetentionCheck."""

from unittest.mock import MagicMock, patch

import pytest

from ai_qa.checks.context_retention import ContextRetentionCheck
from ai_qa.core.context import Context
from ai_qa.exceptions import ContextValidationError


def test_context_retention_check_name():
    """Test ContextRetentionCheck has correct name."""
    check = ContextRetentionCheck()
    assert check.name == "context_retention"


def test_context_retention_check_requires_llm():
    """Test ContextRetentionCheck requires_llm property."""
    check = ContextRetentionCheck()
    assert check.requires_llm is True


@patch("ai_qa.checks.context_retention.get_provider")
def test_context_retention_good_memory(mock_get_provider):
    """Test AI with good memory."""
    check = ContextRetentionCheck()
    ctx = Context(output="test", 
        conversation=[
            {"role": "user", "content": "I love Python"},
            {"role": "assistant", "content": "Python is great"},
            {
                "role": "user",
                "content": "Remember what I said earlier?",
            },
            {"role": "assistant", "content": "Yes, you mentioned Python"},
        ]
    )

    mock_provider = MagicMock()
    mock_provider.judge.return_value = {
        "score": 1.0,
        "evidence": "Perfect memory",
    }
    mock_get_provider.return_value = mock_provider

    result = check.evaluate(ctx)
    assert result.score == 1.0
    assert result.passed


@patch("ai_qa.checks.context_retention.get_provider")
def test_context_retention_poor_memory(mock_get_provider):
    """Test AI with poor memory."""
    check = ContextRetentionCheck(threshold=0.8)
    ctx = Context(output="test", 
        conversation=[
            {"role": "user", "content": "I love Python"},
            {"role": "assistant", "content": "Got it"},
            {"role": "user", "content": "What did I say?"},
            {
                "role": "assistant",
                "content": "I don't remember",
            },
        ]
    )

    mock_provider = MagicMock()
    mock_provider.judge.return_value = {
        "score": 0.2,
        "evidence": "Poor memory",
    }
    mock_get_provider.return_value = mock_provider

    result = check.evaluate(ctx)
    assert result.score == 0.2
    assert not result.passed


def test_context_retention_no_conversation():
    """Test error when conversation missing."""
    check = ContextRetentionCheck()
    ctx = Context(output="test")
    with pytest.raises(ContextValidationError):
        check.evaluate(ctx)


def test_context_retention_single_turn():
    """Test error with single conversation turn."""
    check = ContextRetentionCheck()
    ctx = Context(output="test", 
        conversation=[
            {"role": "user", "content": "Hello"},
        ]
    )
    with pytest.raises(ContextValidationError):
        check.evaluate(ctx)
