"""Declarative MCP workflow DSL for asserting agent behavior patterns."""

from __future__ import annotations

import fnmatch
from typing import TYPE_CHECKING, Any

from pydantic import BaseModel

from ai_qa.core.check import BaseCheck
from ai_qa.core.result import CheckResult
from ai_qa.exceptions import ContextValidationError

if TYPE_CHECKING:
    from ai_qa.core.context import Context


class WorkflowStep(BaseModel):
    """A single step in an expected workflow.

    Attributes:
        tool: Tool name pattern. Supports glob patterns (e.g. "search_*")
            and pipe-separated alternatives (e.g. "create_audience|create_patient_audience").
        required: Whether this step must appear in the trace.
        description: Human-readable description of this step's purpose.
        max_calls: Maximum times this tool pattern can appear.
        args_contain: Keys that must be present in the tool's arguments.
    """

    tool: str
    required: bool = True
    description: str = ""
    max_calls: int | None = None
    args_contain: list[str] | None = None

    def matches(self, tool_name: str) -> bool:
        """Check if a tool name matches this step's pattern.

        Supports:
        - Exact match: "search_diagnosis"
        - Glob: "search_*"
        - Alternatives: "create_audience|create_patient_audience"

        Args:
            tool_name: Actual tool name to test.

        Returns:
            True if the tool name matches.
        """
        if "|" in self.tool:
            return any(
                fnmatch.fnmatch(tool_name, alt.strip())
                for alt in self.tool.split("|")
            )
        return fnmatch.fnmatch(tool_name, self.tool)


class WorkflowConstraints(BaseModel):
    """Global constraints applied to the entire workflow.

    Attributes:
        max_steps: Maximum total tool calls allowed.
        disallowed_tools: Tools that must never be called.
        must_not_repeat: Tools that should only be called once.
        max_retries_per_tool: Maximum times the same tool can be called.
        require_error_recovery: If True, failed tool calls must be
            followed by a retry or alternative approach.
    """

    max_steps: int | None = None
    disallowed_tools: list[str] = []
    must_not_repeat: list[str] = []
    max_retries_per_tool: int | None = None
    require_error_recovery: bool = False


class MCPWorkflow(BaseCheck):
    """Declarative workflow check for MCP agent tool-call sequences.

    Define expected workflows with glob-pattern steps, constraints, and
    use it directly as a check in validate().

    Example:
        workflow = MCPWorkflow(
            workflow_name="audience_creation",
            steps=[
                WorkflowStep(tool="search_*", required=True),
                WorkflowStep(tool="get_*_stats", required=False),
                WorkflowStep(
                    tool="create_audience|create_patient_audience",
                    required=True,
                ),
            ],
            constraints=WorkflowConstraints(
                max_steps=10,
                disallowed_tools=["submit_query"],
                must_not_repeat=["create_audience"],
            ),
        )
        result = validate(output=resp, tool_calls=trace, checks=[workflow])
    """

    name = "mcp_workflow"
    requires_llm = False
    version = "1.0.0"
    description = "Validate agent tool-call sequence against a declared workflow"

    def __init__(
        self,
        workflow_name: str = "unnamed",
        steps: list[WorkflowStep] | list[dict[str, Any]] | None = None,
        constraints: WorkflowConstraints | dict[str, Any] | None = None,
        threshold: float = 0.5,
    ) -> None:
        """Initialize MCPWorkflow.

        Args:
            workflow_name: Human-readable name for this workflow.
            steps: Expected workflow steps (ordered).
            constraints: Global workflow constraints.
            threshold: Score threshold for passing.
        """
        super().__init__(threshold=threshold)
        self.workflow_name = workflow_name

        if steps is None:
            self.steps: list[WorkflowStep] = []
        else:
            self.steps = [
                WorkflowStep(**s) if isinstance(s, dict) else s for s in steps
            ]

        if constraints is None:
            self.constraints = WorkflowConstraints()
        elif isinstance(constraints, dict):
            self.constraints = WorkflowConstraints(**constraints)
        else:
            self.constraints = constraints

    def validate_context(self, ctx: Context) -> None:
        """Validate that tool_calls are present."""
        if ctx.tool_calls is None:
            raise ContextValidationError(
                "tool_calls is required for MCPWorkflow check. "
                "Pass tool_calls as a list of dicts with 'name' and 'arguments' keys, "
                "or use MCPTrace.to_tool_calls() to convert from an MCPTrace."
            )

    def evaluate(self, ctx: Context) -> CheckResult:
        """Evaluate the agent's tool calls against the workflow definition.

        Scoring:
        - Start at 1.0
        - Deduct for each violation found
        - Violations include: missing required steps, constraint breaches,
          disallowed tools, exceeded limits

        Args:
            ctx: Context with tool_calls to evaluate.

        Returns:
            CheckResult with workflow compliance score.
        """
        self.validate_context(ctx)
        assert ctx.tool_calls is not None

        actual_tools = self._extract_tool_names(ctx.tool_calls)
        violations: list[str] = []
        total_checks = 0

        # --- Constraint checks ---

        # Max steps
        if self.constraints.max_steps is not None:
            total_checks += 1
            if len(actual_tools) > self.constraints.max_steps:
                violations.append(
                    f"Exceeded max steps: {len(actual_tools)} > "
                    f"{self.constraints.max_steps}"
                )

        # Disallowed tools
        for tool in actual_tools:
            for disallowed in self.constraints.disallowed_tools:
                total_checks += 1
                if fnmatch.fnmatch(tool, disallowed):
                    violations.append(f"Disallowed tool called: '{tool}'")

        # Must not repeat
        for tool_pattern in self.constraints.must_not_repeat:
            total_checks += 1
            count = sum(
                1 for t in actual_tools if fnmatch.fnmatch(t, tool_pattern)
            )
            if count > 1:
                violations.append(
                    f"Tool '{tool_pattern}' called {count} times "
                    f"but must not repeat"
                )

        # Max retries per tool
        if self.constraints.max_retries_per_tool is not None:
            tool_counts: dict[str, int] = {}
            for t in actual_tools:
                tool_counts[t] = tool_counts.get(t, 0) + 1
            for t, count in tool_counts.items():
                total_checks += 1
                if count > self.constraints.max_retries_per_tool:
                    violations.append(
                        f"Tool '{t}' called {count} times, "
                        f"max retries is {self.constraints.max_retries_per_tool}"
                    )

        # Error recovery
        if self.constraints.require_error_recovery:
            for i, call in enumerate(ctx.tool_calls):
                if isinstance(call, dict) and call.get("error"):
                    total_checks += 1
                    if i == len(ctx.tool_calls) - 1:
                        violations.append(
                            f"Tool '{call.get('name')}' failed with no "
                            f"recovery attempt (was the last call)"
                        )

        # --- Step sequence checks ---
        if self.steps:
            step_violations = self._check_step_sequence(
                actual_tools, ctx.tool_calls
            )
            total_checks += len(self.steps)
            violations.extend(step_violations)

        # --- Scoring ---
        score = (
            1.0
            if total_checks == 0
            else max(0.0, 1.0 - len(violations) / max(total_checks, 1))
        )

        passed = score >= self.threshold and len(violations) == 0
        if violations:
            evidence = (
                f"Workflow '{self.workflow_name}' has "
                f"{len(violations)} violation(s): "
                + "; ".join(violations)
            )
        else:
            evidence = (
                f"Workflow '{self.workflow_name}' completed successfully. "
                f"Tools called: {actual_tools}"
            )

        return CheckResult(
            name=self.name,
            score=score,
            passed=passed,
            threshold=self.threshold,
            evidence=evidence,
            check_version=self.version,
        )

    def _check_step_sequence(
        self,
        actual_tools: list[str],
        tool_calls: list[dict[str, Any]],
    ) -> list[str]:
        """Check that expected steps appear in order (flexible subsequence).

        For optional steps that aren't found, the search index resets so
        subsequent required steps can still match from the same position.

        Args:
            actual_tools: Ordered tool names from the trace.
            tool_calls: Raw tool call dicts (for args_contain checks).

        Returns:
            List of violation messages.
        """
        violations: list[str] = []
        actual_idx = 0

        for step in self.steps:
            found = False
            search_start = actual_idx

            while actual_idx < len(actual_tools):
                if step.matches(actual_tools[actual_idx]):
                    found = True

                    if step.args_contain and actual_idx < len(tool_calls):
                        call = tool_calls[actual_idx]
                        args = call.get("arguments", call.get("params", {}))
                        for key in step.args_contain:
                            if key not in args:
                                violations.append(
                                    f"Step '{step.tool}': missing required "
                                    f"argument '{key}'"
                                )

                    actual_idx += 1
                    break
                actual_idx += 1

            if not found:
                # Reset index so the next step can search from same position
                actual_idx = search_start
                if step.required:
                    desc = f" ({step.description})" if step.description else ""
                    violations.append(
                        f"Required step '{step.tool}'{desc} not found "
                        f"in tool sequence"
                    )

            if step.max_calls is not None:
                count = sum(
                    1 for t in actual_tools if step.matches(t)
                )
                if count > step.max_calls:
                    violations.append(
                        f"Step '{step.tool}' called {count} times, "
                        f"max is {step.max_calls}"
                    )

        return violations

    @staticmethod
    def _extract_tool_names(tool_calls: list[dict[str, Any]]) -> list[str]:
        """Extract ordered tool names from Context.tool_calls format."""
        names: list[str] = []
        for call in tool_calls:
            if isinstance(call, dict):
                name = call.get("name", call.get("tool", ""))
                if name:
                    names.append(name)
        return names
