"""TrajectoryCheck: Validate sequence of tool calls in agent workflow."""

from __future__ import annotations

from typing import TYPE_CHECKING, Any

from ai_qa.checks.registry import register
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


@register("trajectory")
class TrajectoryCheck(BaseCheck):
    """Check that an agent followed the expected workflow trajectory.

    Args:
        expected_trajectory: List of dicts with "tool" and optional "order" keys.
        match: "strict" for exact order, "flexible" for subsequence.
        max_steps: Maximum allowed tool calls (optional).
        disallowed_tools: Tools that should never be called.
    """

    requires_llm = False
    version = "1.0.0"
    description = "Validate sequence of tool calls in agent workflow"

    def __init__(
        self,
        expected_trajectory: list[dict[str, Any]],
        match: str = "flexible",
        max_steps: int | None = None,
        disallowed_tools: list[str] | None = None,
        threshold: float = 0.5,
    ) -> None:
        """Initialize TrajectoryCheck.

        Args:
            expected_trajectory: Expected tool sequence.
            match: "strict" or "flexible".
            max_steps: Maximum allowed steps.
            disallowed_tools: Tools to forbid.
            threshold: Score threshold for passing.

        Raises:
            ValueError: If match is invalid.
        """
        if match not in ("strict", "flexible"):
            msg = f"match must be 'strict' or 'flexible', got {match}"
            raise ValueError(msg)
        super().__init__(threshold=threshold)
        self.expected_trajectory = expected_trajectory
        self.match = match
        self.max_steps = max_steps
        self.disallowed_tools = disallowed_tools or []

    def validate_context(self, ctx: Context) -> None:
        """Validate that tool_calls are available.

        Args:
            ctx: The context to validate.

        Raises:
            ContextValidationError: If tool_calls is missing.
        """
        if ctx.tool_calls is None:
            msg = "tool_calls is required for TrajectoryCheck"
            raise ContextValidationError(msg)

    def evaluate(self, ctx: Context) -> CheckResult:
        """Evaluate tool call trajectory.

        Args:
            ctx: Context with tool_calls to evaluate.

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

        # Extract tool names in order
        actual_tools = []
        for call in ctx.tool_calls:
            if isinstance(call, dict):
                tool = call.get("name")
                if tool:
                    actual_tools.append(tool)

        # Check max_steps constraint
        if self.max_steps is not None and len(actual_tools) > self.max_steps:
            return CheckResult(
                name=self.name,
                score=0.0,
                passed=False,
                threshold=self.threshold,
                evidence=(
                    f"Exceeded max steps: {len(actual_tools)} > "
                    f"{self.max_steps}"
                ),
                cost_usd=0.0,
                tokens_used={},
                check_version=self.version,
            )

        # Check disallowed_tools
        for tool in actual_tools:
            if tool in self.disallowed_tools:
                return CheckResult(
                    name=self.name,
                    score=0.0,
                    passed=False,
                    threshold=self.threshold,
                    evidence=f"Disallowed tool called: '{tool}'",
                    cost_usd=0.0,
                    tokens_used={},
                    check_version=self.version,
                )

        # Extract expected tool names
        expected_tools = [t.get("tool") for t in self.expected_trajectory if "tool" in t]

        # Check trajectory match
        if self.match == "strict":
            match = actual_tools == expected_tools
        else:  # flexible (subsequence)
            # Check if expected_tools is a subsequence of actual_tools
            match = self._is_subsequence(expected_tools, actual_tools)

        score = 1.0 if match else 0.5
        evidence = (
            f"Trajectory matches: {actual_tools}"
            if match
            else (
                f"Trajectory mismatch: expected "
                f"{expected_tools}, got {actual_tools}"
            )
        )

        return CheckResult(
            name=self.name,
            score=score,
            passed=score >= self.threshold,
            threshold=self.threshold,
            evidence=evidence,
            cost_usd=0.0,
            tokens_used={},
            check_version=self.version,
        )

    @staticmethod
    def _is_subsequence(expected: list[Any], actual: list[Any]) -> bool:
        """Check if expected is a subsequence of actual.

        Args:
            expected: Expected sequence.
            actual: Actual sequence.

        Returns:
            True if expected is a subsequence of actual.
        """
        if not expected:
            return True
        exp_idx = 0
        for item in actual:
            if item == expected[exp_idx]:
                exp_idx += 1
                if exp_idx == len(expected):
                    return True
        return False
