"""Tests for MCP-specific checks."""

from __future__ import annotations

from unittest.mock import MagicMock, patch

import pytest

from ai_qa.core.context import Context
from ai_qa.exceptions import ContextValidationError
from ai_qa.mcp.checks import (
    MCPToolEfficiencyCheck,
    MCPToolErrorHandlingCheck,
    MCPToolParamsSchemaCheck,
    MCPToolResultUsageCheck,
)


class TestMCPToolParamsSchemaCheck:
    def _make_ctx(self, tool_calls: list[dict]) -> Context:
        return Context(output="test", tool_calls=tool_calls)

    def test_valid_params(self) -> None:
        check = MCPToolParamsSchemaCheck(
            tool_schemas={
                "search_diagnosis": {
                    "type": "object",
                    "properties": {"text": {"type": "string"}},
                    "required": ["text"],
                }
            }
        )
        ctx = self._make_ctx([
            {"name": "search_diagnosis", "arguments": {"text": "diabetes"}},
        ])
        result = check.evaluate(ctx)
        assert result.passed is True
        assert result.score == 1.0

    def test_invalid_params_missing_required(self) -> None:
        check = MCPToolParamsSchemaCheck(
            tool_schemas={
                "search_diagnosis": {
                    "type": "object",
                    "properties": {"text": {"type": "string"}},
                    "required": ["text"],
                }
            }
        )
        ctx = self._make_ctx([
            {"name": "search_diagnosis", "arguments": {}},
        ])
        result = check.evaluate(ctx)
        assert result.passed is False
        assert result.score == 0.0
        assert "search_diagnosis" in result.evidence

    def test_invalid_params_wrong_type(self) -> None:
        check = MCPToolParamsSchemaCheck(
            tool_schemas={
                "search_diagnosis": {
                    "type": "object",
                    "properties": {"text": {"type": "string"}},
                    "required": ["text"],
                }
            }
        )
        ctx = self._make_ctx([
            {"name": "search_diagnosis", "arguments": {"text": 123}},
        ])
        result = check.evaluate(ctx)
        assert result.passed is False

    def test_multiple_tools_partial_failure(self) -> None:
        check = MCPToolParamsSchemaCheck(
            tool_schemas={
                "search_diagnosis": {
                    "type": "object",
                    "properties": {"text": {"type": "string"}},
                    "required": ["text"],
                },
                "search_drug": {
                    "type": "object",
                    "properties": {"text": {"type": "string"}},
                    "required": ["text"],
                },
            }
        )
        ctx = self._make_ctx([
            {"name": "search_diagnosis", "arguments": {"text": "diabetes"}},
            {"name": "search_drug", "arguments": {}},
        ])
        result = check.evaluate(ctx)
        assert result.score == 0.5

    def test_unmatched_tools_skipped(self) -> None:
        check = MCPToolParamsSchemaCheck(
            tool_schemas={
                "search_diagnosis": {
                    "type": "object",
                    "properties": {"text": {"type": "string"}},
                }
            }
        )
        ctx = self._make_ctx([
            {"name": "create_audience", "arguments": {"filter": "{}"}},
        ])
        result = check.evaluate(ctx)
        assert result.passed is True
        assert "nothing to validate" in result.evidence.lower()

    def test_no_tool_calls_raises(self) -> None:
        check = MCPToolParamsSchemaCheck(tool_schemas={})
        ctx = Context(output="test")
        with pytest.raises(ContextValidationError):
            check.evaluate(ctx)

    def test_registered_name(self) -> None:
        from ai_qa.checks import get_check

        check = get_check(
            "mcp_tool_params_schema",
            tool_schemas={"test": {"type": "object"}},
        )
        assert check.name == "mcp_tool_params_schema"


class TestMCPToolEfficiencyCheck:
    def _make_ctx(self, tool_calls: list[dict]) -> Context:
        return Context(output="test", tool_calls=tool_calls)

    def test_no_duplicates(self) -> None:
        check = MCPToolEfficiencyCheck()
        ctx = self._make_ctx([
            {"name": "search_diagnosis", "arguments": {"text": "diabetes"}},
            {"name": "search_drug", "arguments": {"text": "metformin"}},
        ])
        result = check.evaluate(ctx)
        assert result.passed is True
        assert result.score == 1.0
        assert "No duplicate" in result.evidence

    def test_exact_duplicate_detected(self) -> None:
        check = MCPToolEfficiencyCheck()
        ctx = self._make_ctx([
            {"name": "search_diagnosis", "arguments": {"text": "diabetes"}},
            {"name": "search_diagnosis", "arguments": {"text": "diabetes"}},
        ])
        result = check.evaluate(ctx)
        assert result.score == 0.5
        assert "Duplicate" in result.evidence

    def test_same_tool_different_args_is_ok(self) -> None:
        check = MCPToolEfficiencyCheck()
        ctx = self._make_ctx([
            {"name": "search_diagnosis", "arguments": {"text": "diabetes"}},
            {"name": "search_diagnosis", "arguments": {"text": "asthma"}},
        ])
        result = check.evaluate(ctx)
        assert result.score == 1.0

    def test_multiple_duplicates(self) -> None:
        check = MCPToolEfficiencyCheck()
        ctx = self._make_ctx([
            {"name": "a", "arguments": {"x": 1}},
            {"name": "a", "arguments": {"x": 1}},
            {"name": "a", "arguments": {"x": 1}},
            {"name": "b", "arguments": {}},
        ])
        result = check.evaluate(ctx)
        assert result.score == 0.5  # 2 duplicates out of 4

    def test_empty_calls(self) -> None:
        check = MCPToolEfficiencyCheck()
        ctx = self._make_ctx([])
        result = check.evaluate(ctx)
        assert result.passed is True

    def test_no_tool_calls_raises(self) -> None:
        check = MCPToolEfficiencyCheck()
        ctx = Context(output="test")
        with pytest.raises(ContextValidationError):
            check.evaluate(ctx)

    def test_registered_name(self) -> None:
        from ai_qa.checks import get_check

        check = get_check("mcp_tool_efficiency")
        assert check.name == "mcp_tool_efficiency"


class TestMCPToolErrorHandlingCheck:
    def _make_ctx(self, tool_calls: list[dict]) -> Context:
        return Context(output="test", tool_calls=tool_calls)

    def test_no_errors(self) -> None:
        check = MCPToolErrorHandlingCheck()
        ctx = self._make_ctx([
            {"name": "search_diagnosis", "arguments": {}},
            {"name": "create_audience", "arguments": {}},
        ])
        result = check.evaluate(ctx)
        assert result.passed is True
        assert result.score == 1.0
        assert "No tool errors" in result.evidence

    def test_error_with_retry(self) -> None:
        check = MCPToolErrorHandlingCheck()
        ctx = self._make_ctx([
            {"name": "search_diagnosis", "arguments": {}, "error": "timeout"},
            {"name": "search_diagnosis", "arguments": {}},
            {"name": "create_audience", "arguments": {}},
        ])
        result = check.evaluate(ctx)
        assert result.passed is True
        assert "recovered" in result.evidence.lower()

    def test_error_without_retry(self) -> None:
        check = MCPToolErrorHandlingCheck()
        ctx = self._make_ctx([
            {"name": "search_diagnosis", "arguments": {}, "error": "timeout"},
        ])
        result = check.evaluate(ctx)
        assert result.passed is False
        assert "no recovery" in result.evidence.lower()

    def test_error_with_alternative_tool(self) -> None:
        check = MCPToolErrorHandlingCheck()
        ctx = self._make_ctx([
            {"name": "search_diagnosis", "arguments": {}, "error": "timeout"},
            {"name": "search_drug", "arguments": {}},
        ])
        result = check.evaluate(ctx)
        assert result.passed is True
        assert result.score == 1.0

    def test_excessive_retries(self) -> None:
        check = MCPToolErrorHandlingCheck(max_retries=2)
        ctx = self._make_ctx([
            {"name": "search_diagnosis", "arguments": {}, "error": "fail"},
            {"name": "search_diagnosis", "arguments": {}},
            {"name": "search_diagnosis", "arguments": {}},
            {"name": "search_diagnosis", "arguments": {}},
        ])
        result = check.evaluate(ctx)
        assert result.passed is False
        assert "max" in result.evidence.lower()

    def test_success_false_treated_as_error(self) -> None:
        check = MCPToolErrorHandlingCheck()
        ctx = self._make_ctx([
            {
                "name": "create_audience",
                "arguments": {},
                "result": {"success": False, "error": "Invalid filter"},
            },
        ])
        result = check.evaluate(ctx)
        assert result.passed is False

    def test_no_tool_calls_raises(self) -> None:
        check = MCPToolErrorHandlingCheck()
        ctx = Context(output="test")
        with pytest.raises(ContextValidationError):
            check.evaluate(ctx)

    def test_registered_name(self) -> None:
        from ai_qa.checks import get_check

        check = get_check("mcp_tool_error_handling")
        assert check.name == "mcp_tool_error_handling"


class TestMCPToolResultUsageCheck:
    def test_requires_llm(self) -> None:
        check = MCPToolResultUsageCheck()
        assert check.requires_llm is True

    def test_no_tool_calls_raises(self) -> None:
        check = MCPToolResultUsageCheck()
        ctx = Context(output="test")
        with pytest.raises(ContextValidationError):
            check.evaluate(ctx)

    @patch("ai_qa.providers.get_provider")
    def test_evaluate_with_mock_provider(
        self, mock_get_provider: MagicMock
    ) -> None:
        mock_provider = MagicMock()
        mock_provider.judge.return_value = {
            "score": 0.9,
            "evidence": "Agent used tool results appropriately",
            "cost_usd": 0.001,
            "tokens": {"input": 100, "output": 50},
        }
        mock_get_provider.return_value = mock_provider

        check = MCPToolResultUsageCheck()
        ctx = Context(
            input="Find oncologists in NY",
            output="Found 500 oncologists in New York",
            tool_calls=[
                {
                    "name": "search_hcp_specialities",
                    "arguments": {"text": "oncology"},
                    "result": {"success": True, "count": 2},
                },
                {
                    "name": "get_hcp_stats",
                    "arguments": {"dimension_type": "GEO"},
                    "result": {"success": True, "data": [{"name": "NY", "y": 500}]},
                },
            ],
        )
        result = check.evaluate(ctx)
        assert result.passed is True
        assert result.score == 0.9
        mock_provider.judge.assert_called_once()

    def test_registered_name(self) -> None:
        from ai_qa.checks import get_check

        check = get_check("mcp_tool_result_usage")
        assert check.name == "mcp_tool_result_usage"
