"""ToolUseCheck: Validate agent tool selection and parameters."""

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
from ai_qa.providers import get_provider

if TYPE_CHECKING:
    from ai_qa.core.context import Context

JUDGE_PROMPT = """You are an expert at evaluating whether two sets of parameters
are semantically equivalent for the same function call.

Given:
1. Expected parameters
2. Actual parameters provided

Compare them and determine if they match semantically (same intent and values,
possibly with different keys/structure). Score higher if they match closely.

Respond with a JSON object:
{
  "score": <float 0.0-1.0, where 1.0 means parameters are equivalent>,
  "evidence": "<explanation of parameter match/mismatch>"
}
"""


@register("tool_use")
class ToolUseCheck(BaseCheck):
    """Check that an agent called the correct tool with correct parameters.

    Args:
        expected_tool: Name of the tool that should be called.
        expected_params: Expected parameters for the tool call (optional).
        param_match: How to match parameters: "exact", "subset", or "semantic".
    """

    requires_llm = False
    version = "1.0.0"
    description = "Validate agent tool selection and parameters"

    def __init__(
        self,
        expected_tool: str,
        expected_params: dict[str, Any] | None = None,
        param_match: str = "exact",
        threshold: float = 0.5,
    ) -> None:
        """Initialize ToolUseCheck.

        Args:
            expected_tool: Tool name to expect.
            expected_params: Expected parameters dict.
            param_match: "exact", "subset", or "semantic".
            threshold: Score threshold for passing.

        Raises:
            ValueError: If param_match is invalid.
        """
        if param_match not in ("exact", "subset", "semantic"):
            msg = (
                f"param_match must be 'exact', 'subset', or 'semantic', "
                f"got {param_match}"
            )
            raise ValueError(msg)
        super().__init__(threshold=threshold)
        self.expected_tool = expected_tool
        self.expected_params = expected_params or {}
        self.param_match = param_match

    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 ToolUseCheck"
            raise ContextValidationError(msg)

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

        Args:
            ctx: Context with tool_calls to evaluate.

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

        # Find the expected tool in tool_calls
        tool_found = False
        tool_call = None
        for call in ctx.tool_calls:
            if isinstance(call, dict) and call.get("name") == self.expected_tool:
                tool_found = True
                tool_call = call
                break

        if not tool_found:
            return CheckResult(
                name=self.name,
                score=0.0,
                passed=False,
                threshold=self.threshold,
                evidence=f"Expected tool '{self.expected_tool}' not found",
                cost_usd=0.0,
                tokens_used={},
                check_version=self.version,
            )

        # If no params expected, tool match alone is success
        if not self.expected_params:
            return CheckResult(
                name=self.name,
                score=1.0,
                passed=True,
                threshold=self.threshold,
                evidence=f"Tool '{self.expected_tool}' called correctly",
                cost_usd=0.0,
                tokens_used={},
                check_version=self.version,
            )

        # Compare parameters
        assert tool_call is not None
        actual_params = tool_call.get("params", {})

        if self.param_match == "exact":
            match = actual_params == self.expected_params
            score = 1.0 if match else 0.5
            evidence = (
                "Parameters match exactly"
                if match
                else (
                    f"Parameters mismatch: expected "
                    f"{self.expected_params}, got {actual_params}"
                )
            )
        elif self.param_match == "subset":
            # Check if expected_params is a subset of actual_params
            match = all(
                actual_params.get(k) == v for k, v in self.expected_params.items()
            )
            score = 1.0 if match else 0.5
            evidence = (
                "Expected parameters are subset of actual"
                if match
                else f"Expected parameters not found: {self.expected_params}"
            )
        else:  # semantic
            # Use LLM to judge semantic equivalence
            provider = get_provider()
            prompt = (
                f"EXPECTED:\n{self.expected_params}\n\n"
                f"ACTUAL:\n{actual_params}\n\n"
                "Are these semantically equivalent?"
            )
            resp = provider.judge(system=JUDGE_PROMPT, user=prompt)
            score = resp.get("score", 0.5)
            evidence = resp.get("evidence", "Semantic comparison performed")
            return CheckResult(
                name=self.name,
                score=score,
                passed=score >= self.threshold,
                threshold=self.threshold,
                evidence=evidence,
                cost_usd=resp.get("cost_usd", 0.0),
                tokens_used=resp.get("tokens", {}),
                check_version=self.version,
            )

        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,
        )
