"""Tests for MCPWorkflow DSL and workflow check."""

from __future__ import annotations

import pytest

from ai_qa.core.context import Context
from ai_qa.exceptions import ContextValidationError
from ai_qa.mcp.workflow import MCPWorkflow, WorkflowConstraints, WorkflowStep


class TestWorkflowStep:
    def test_exact_match(self) -> None:
        step = WorkflowStep(tool="search_diagnosis")
        assert step.matches("search_diagnosis") is True
        assert step.matches("search_drug") is False

    def test_glob_match(self) -> None:
        step = WorkflowStep(tool="search_*")
        assert step.matches("search_diagnosis") is True
        assert step.matches("search_drug") is True
        assert step.matches("create_audience") is False

    def test_pipe_alternatives(self) -> None:
        step = WorkflowStep(tool="create_audience|create_patient_audience")
        assert step.matches("create_audience") is True
        assert step.matches("create_patient_audience") is True
        assert step.matches("search_diagnosis") is False

    def test_glob_in_alternatives(self) -> None:
        step = WorkflowStep(tool="search_*|get_*_stats")
        assert step.matches("search_diagnosis") is True
        assert step.matches("get_hcp_stats") is True
        assert step.matches("create_audience") is False


class TestMCPWorkflow:
    def _make_ctx(
        self, tool_names: list[str], **kwargs: object
    ) -> Context:
        tool_calls = [
            {"name": name, "arguments": {}} for name in tool_names
        ]
        return Context(
            output="test output",
            tool_calls=tool_calls,
            **kwargs,
        )

    def test_empty_workflow_passes(self) -> None:
        workflow = MCPWorkflow(workflow_name="empty")
        ctx = self._make_ctx(["search_diagnosis", "create_audience"])
        result = workflow.evaluate(ctx)
        assert result.passed is True
        assert result.score == 1.0

    def test_required_step_present(self) -> None:
        workflow = MCPWorkflow(
            workflow_name="basic",
            steps=[
                WorkflowStep(tool="search_diagnosis", required=True),
                WorkflowStep(tool="create_audience", required=True),
            ],
        )
        ctx = self._make_ctx(["search_diagnosis", "get_hcp_stats", "create_audience"])
        result = workflow.evaluate(ctx)
        assert result.passed is True

    def test_required_step_missing(self) -> None:
        workflow = MCPWorkflow(
            workflow_name="missing_step",
            steps=[
                WorkflowStep(tool="search_diagnosis", required=True),
                WorkflowStep(tool="create_audience", required=True),
            ],
        )
        ctx = self._make_ctx(["search_diagnosis", "get_hcp_stats"])
        result = workflow.evaluate(ctx)
        assert result.passed is False
        assert "create_audience" in result.evidence

    def test_optional_step_missing_is_ok(self) -> None:
        workflow = MCPWorkflow(
            workflow_name="optional",
            steps=[
                WorkflowStep(tool="search_*", required=True),
                WorkflowStep(tool="get_*_stats", required=False),
                WorkflowStep(tool="create_audience", required=True),
            ],
        )
        ctx = self._make_ctx(["search_diagnosis", "create_audience"])
        result = workflow.evaluate(ctx)
        assert result.passed is True

    def test_glob_step_matching(self) -> None:
        workflow = MCPWorkflow(
            workflow_name="glob",
            steps=[
                WorkflowStep(tool="search_*", required=True),
            ],
        )
        ctx = self._make_ctx(["search_drug"])
        result = workflow.evaluate(ctx)
        assert result.passed is True

    def test_max_steps_constraint(self) -> None:
        workflow = MCPWorkflow(
            workflow_name="max_steps",
            constraints=WorkflowConstraints(max_steps=3),
        )
        ctx = self._make_ctx(["a", "b", "c", "d", "e"])
        result = workflow.evaluate(ctx)
        assert result.passed is False
        assert "Exceeded max steps" in result.evidence

    def test_disallowed_tools(self) -> None:
        workflow = MCPWorkflow(
            workflow_name="disallowed",
            constraints=WorkflowConstraints(
                disallowed_tools=["submit_query"]
            ),
        )
        ctx = self._make_ctx(["search_diagnosis", "submit_query"])
        result = workflow.evaluate(ctx)
        assert result.passed is False
        assert "Disallowed tool" in result.evidence

    def test_must_not_repeat(self) -> None:
        workflow = MCPWorkflow(
            workflow_name="no_repeat",
            constraints=WorkflowConstraints(
                must_not_repeat=["create_audience"]
            ),
        )
        ctx = self._make_ctx([
            "search_diagnosis",
            "create_audience",
            "create_audience",
        ])
        result = workflow.evaluate(ctx)
        assert result.passed is False
        assert "must not repeat" in result.evidence

    def test_max_retries_per_tool(self) -> None:
        workflow = MCPWorkflow(
            workflow_name="retries",
            constraints=WorkflowConstraints(max_retries_per_tool=2),
        )
        ctx = self._make_ctx(["search_drug", "search_drug", "search_drug"])
        result = workflow.evaluate(ctx)
        assert result.passed is False

    def test_error_recovery_required(self) -> None:
        workflow = MCPWorkflow(
            workflow_name="recovery",
            constraints=WorkflowConstraints(require_error_recovery=True),
        )
        ctx = Context(
            output="test",
            tool_calls=[
                {"name": "search_diagnosis", "arguments": {}, "error": "timeout"},
            ],
        )
        result = workflow.evaluate(ctx)
        assert result.passed is False
        assert "no recovery" in result.evidence.lower()

    def test_args_contain_validation(self) -> None:
        workflow = MCPWorkflow(
            workflow_name="args_check",
            steps=[
                WorkflowStep(
                    tool="search_diagnosis",
                    required=True,
                    args_contain=["text"],
                ),
            ],
        )
        ctx = Context(
            output="test",
            tool_calls=[
                {"name": "search_diagnosis", "arguments": {"text": "diabetes"}},
            ],
        )
        result = workflow.evaluate(ctx)
        assert result.passed is True

    def test_args_contain_missing(self) -> None:
        workflow = MCPWorkflow(
            workflow_name="args_missing",
            steps=[
                WorkflowStep(
                    tool="search_diagnosis",
                    required=True,
                    args_contain=["text"],
                ),
            ],
        )
        ctx = Context(
            output="test",
            tool_calls=[
                {"name": "search_diagnosis", "arguments": {}},
            ],
        )
        result = workflow.evaluate(ctx)
        assert result.passed is False
        assert "missing required argument" in result.evidence

    def test_no_tool_calls_raises(self) -> None:
        workflow = MCPWorkflow(workflow_name="no_calls")
        ctx = Context(output="test")
        with pytest.raises(ContextValidationError):
            workflow.evaluate(ctx)

    def test_dict_step_initialization(self) -> None:
        workflow = MCPWorkflow(
            workflow_name="dict_init",
            steps=[
                {"tool": "search_*", "required": True},
                {"tool": "create_audience", "required": True},
            ],
        )
        assert len(workflow.steps) == 2
        assert isinstance(workflow.steps[0], WorkflowStep)

    def test_dict_constraints_initialization(self) -> None:
        workflow = MCPWorkflow(
            workflow_name="dict_constraints",
            constraints={"max_steps": 5, "disallowed_tools": ["bad_tool"]},
        )
        assert workflow.constraints.max_steps == 5
        assert "bad_tool" in workflow.constraints.disallowed_tools

    def test_step_max_calls(self) -> None:
        workflow = MCPWorkflow(
            workflow_name="step_max",
            steps=[
                WorkflowStep(tool="search_*", max_calls=2),
            ],
        )
        ctx = self._make_ctx([
            "search_diagnosis",
            "search_drug",
            "search_procedure",
        ])
        result = workflow.evaluate(ctx)
        assert result.passed is False
        assert "max is 2" in result.evidence

    def test_complex_deepintent_workflow(self) -> None:
        """Simulate a real DeepIntent audience creation workflow."""
        workflow = MCPWorkflow(
            workflow_name="audience_creation",
            steps=[
                WorkflowStep(
                    tool="search_*",
                    required=True,
                    description="Look up medical codes",
                ),
                WorkflowStep(
                    tool="get_hcp_stats|get_patient_stats",
                    required=False,
                    description="Optional audience sizing",
                ),
                WorkflowStep(
                    tool="create_audience|create_patient_audience",
                    required=True,
                    description="Create the audience",
                    max_calls=1,
                ),
            ],
            constraints=WorkflowConstraints(
                max_steps=10,
                disallowed_tools=["submit_query"],
                must_not_repeat=["create_audience", "create_patient_audience"],
            ),
        )
        ctx = self._make_ctx([
            "search_hcp_specialities",
            "search_diagnosis",
            "get_hcp_stats",
            "create_audience",
        ])
        result = workflow.evaluate(ctx)
        assert result.passed is True
        assert "successfully" in result.evidence
