"""MCP-specific checks for validating agent tool-call behavior.

These checks go beyond basic tool_use and trajectory checks to validate
higher-level patterns specific to MCP agent workflows.
"""

from __future__ import annotations

import json
from collections import Counter
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


TOOL_RESULT_USAGE_PROMPT = """You are an expert at evaluating whether an AI agent
properly used the results from its tool calls to formulate its final response.

Given:
1. The user's original question/request
2. The tool calls the agent made and their results
3. The agent's final response to the user

Evaluate whether:
- The agent actually used the relevant data from tool results
- The agent did not fabricate information that wasn't in the tool results
- Important data from tool results was not ignored
- The final response is consistent with what the tools returned

Respond with ONLY a JSON object:
{
  "score": <float 0.0 to 1.0>,
  "evidence": "<explanation of how well tool results were used>",
  "unused_results": ["<tool results that were ignored>"],
  "fabricated_claims": ["<claims not supported by tool results>"]
}
"""

TOOL_EFFICIENCY_PROMPT = """You are an expert at evaluating the efficiency of
an AI agent's tool-calling strategy.

Given:
1. The user's original question/request
2. The sequence of tool calls the agent made

Evaluate whether:
- Any tool calls were redundant (same tool called with same/similar args)
- Any tool calls were unnecessary for answering the question
- The agent could have answered with fewer tool calls
- The ordering of tool calls was logical

Respond with ONLY a JSON object:
{
  "score": <float 0.0 to 1.0>,
  "evidence": "<explanation of efficiency>",
  "redundant_calls": ["<description of redundant calls>"],
  "unnecessary_calls": ["<description of unnecessary calls>"],
  "suggested_optimization": "<how the workflow could be more efficient>"
}
"""


@register("mcp_tool_params_schema")
class MCPToolParamsSchemaCheck(BaseCheck):
    """Validate that tool call arguments match expected JSON schemas.

    For each tool call in the trace, checks that the arguments conform to
    a declared schema. Useful for catching malformed filter JSON, missing
    required fields, wrong types, etc.

    Args:
        tool_schemas: Mapping of tool name to JSON Schema dict defining
            the expected argument structure. Tools not in this mapping
            are skipped.
    """

    requires_llm = False
    version = "1.0.0"
    description = "Validate tool call arguments against JSON schemas"

    def __init__(
        self,
        tool_schemas: dict[str, dict[str, Any]],
        threshold: float = 1.0,
    ) -> None:
        """Initialize MCPToolParamsSchemaCheck.

        Args:
            tool_schemas: {tool_name: json_schema_dict} for validation.
            threshold: Score threshold for passing.
        """
        super().__init__(threshold=threshold)
        self.tool_schemas = tool_schemas

    def validate_context(self, ctx: Context) -> None:
        if ctx.tool_calls is None:
            raise ContextValidationError(
                "tool_calls is required for MCPToolParamsSchemaCheck"
            )

    def evaluate(self, ctx: Context) -> CheckResult:
        """Validate each tool call's arguments against its schema.

        Args:
            ctx: Context with tool_calls.

        Returns:
            CheckResult — score is proportion of schema-valid calls.
        """
        self.validate_context(ctx)
        assert ctx.tool_calls is not None

        import jsonschema

        validated = 0
        errors: list[str] = []

        for call in ctx.tool_calls:
            if not isinstance(call, dict):
                continue
            tool_name = call.get("name", call.get("tool", ""))
            if tool_name not in self.tool_schemas:
                continue

            validated += 1
            schema = self.tool_schemas[tool_name]
            args = call.get("arguments", call.get("params", {}))

            try:
                jsonschema.validate(args, schema)
            except jsonschema.ValidationError as e:
                errors.append(
                    f"{tool_name}: {e.message} "
                    f"(path: {'.'.join(str(p) for p in e.path)})"
                    if e.path
                    else f"{tool_name}: {e.message}"
                )

        if validated == 0:
            return CheckResult(
                name=self.name,
                score=1.0,
                passed=True,
                threshold=self.threshold,
                evidence="No tool calls matched provided schemas — nothing to validate",
                check_version=self.version,
            )

        score = (validated - len(errors)) / validated
        passed = score >= self.threshold
        evidence = (
            f"All {validated} tool calls have valid arguments"
            if not errors
            else f"{len(errors)} schema violation(s): " + "; ".join(errors)
        )

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


@register("mcp_tool_result_usage")
class MCPToolResultUsageCheck(BaseCheck):
    """Check that the agent actually used tool results in its final response.

    Uses an LLM judge to evaluate whether the agent's output is grounded
    in the data returned by its tool calls, rather than fabricated.
    """

    requires_llm = True
    version = "1.0.0"
    description = "Validate that tool results were used in the agent's response"

    def validate_context(self, ctx: Context) -> None:
        if ctx.tool_calls is None:
            raise ContextValidationError(
                "tool_calls is required for MCPToolResultUsageCheck"
            )

    def evaluate(self, ctx: Context) -> CheckResult:
        """Evaluate tool result usage via LLM judge.

        Args:
            ctx: Context with tool_calls and output.

        Returns:
            CheckResult scoring how well tool results were used.
        """
        self.validate_context(ctx)
        assert ctx.tool_calls is not None

        from ai_qa.providers import get_provider

        provider = get_provider()

        tool_summary = self._summarize_tool_calls(ctx.tool_calls)
        user_prompt = (
            f"USER REQUEST: {ctx.input or '(not provided)'}\n\n"
            f"TOOL CALLS AND RESULTS:\n{tool_summary}\n\n"
            f"AGENT FINAL RESPONSE:\n{ctx.output}"
        )

        resp = provider.judge(system=TOOL_RESULT_USAGE_PROMPT, user=user_prompt)
        score = resp.get("score", 0.5)

        return CheckResult(
            name=self.name,
            score=score,
            passed=score >= self.threshold,
            threshold=self.threshold,
            evidence=resp.get("evidence", "Tool result usage evaluated"),
            cost_usd=resp.get("cost_usd", 0.0),
            tokens_used=resp.get("tokens", {}),
            raw=resp,
            check_version=self.version,
        )

    @staticmethod
    def _summarize_tool_calls(
        tool_calls: list[dict[str, Any]],
        max_result_chars: int = 500,
    ) -> str:
        """Build a human-readable summary of tool calls for the judge prompt."""
        lines: list[str] = []
        for i, call in enumerate(tool_calls, 1):
            name = call.get("name", call.get("tool", "unknown"))
            args = call.get("arguments", call.get("params", {}))
            result = call.get("result")
            error = call.get("error")

            lines.append(f"Call {i}: {name}")
            lines.append(f"  Args: {json.dumps(args, default=str)[:300]}")
            if error:
                lines.append(f"  Error: {error}")
            elif result is not None:
                result_str = json.dumps(result, default=str)
                if len(result_str) > max_result_chars:
                    result_str = result_str[:max_result_chars] + "..."
                lines.append(f"  Result: {result_str}")
            else:
                lines.append("  Result: (none)")
        return "\n".join(lines)


@register("mcp_tool_efficiency")
class MCPToolEfficiencyCheck(BaseCheck):
    """Check whether the agent's tool-calling strategy was efficient.

    Can run in deterministic mode (detecting exact duplicate calls) or
    LLM mode (judging semantic redundancy and unnecessary calls).

    Args:
        use_llm: If True, uses LLM for deeper analysis. If False,
            only checks for exact duplicate calls.
    """

    requires_llm = False
    version = "1.0.0"
    description = "Evaluate efficiency of agent tool-calling strategy"

    def __init__(
        self,
        use_llm: bool = False,
        threshold: float = 0.5,
    ) -> None:
        super().__init__(threshold=threshold)
        self.use_llm = use_llm
        if use_llm:
            self.requires_llm = True

    def validate_context(self, ctx: Context) -> None:
        if ctx.tool_calls is None:
            raise ContextValidationError(
                "tool_calls is required for MCPToolEfficiencyCheck"
            )

    def evaluate(self, ctx: Context) -> CheckResult:
        self.validate_context(ctx)
        assert ctx.tool_calls is not None

        if self.use_llm:
            return self._evaluate_with_llm(ctx)
        return self._evaluate_deterministic(ctx)

    def _evaluate_deterministic(self, ctx: Context) -> CheckResult:
        """Detect exact duplicate tool calls (same tool + same args)."""
        assert ctx.tool_calls is not None

        call_signatures: list[str] = []
        duplicates: list[str] = []

        for call in ctx.tool_calls:
            if not isinstance(call, dict):
                continue
            name = call.get("name", call.get("tool", ""))
            args = call.get("arguments", call.get("params", {}))
            sig = f"{name}:{json.dumps(args, sort_keys=True, default=str)}"

            if sig in call_signatures:
                duplicates.append(name)
            call_signatures.append(sig)

        total = len(call_signatures)
        if total == 0:
            return CheckResult(
                name=self.name,
                score=1.0,
                passed=True,
                threshold=self.threshold,
                evidence="No tool calls to evaluate",
                check_version=self.version,
            )

        score = (total - len(duplicates)) / total
        if duplicates:
            counts = Counter(duplicates)
            dup_desc = ", ".join(
                f"{name} ({c} extra)" for name, c in counts.items()
            )
            evidence = f"Duplicate tool calls detected: {dup_desc}"
        else:
            evidence = f"No duplicate tool calls among {total} calls"

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

    def _evaluate_with_llm(self, ctx: Context) -> CheckResult:
        """Use LLM to evaluate deeper efficiency patterns."""
        assert ctx.tool_calls is not None

        from ai_qa.providers import get_provider

        provider = get_provider()

        tool_summary = MCPToolResultUsageCheck._summarize_tool_calls(
            ctx.tool_calls
        )
        user_prompt = (
            f"USER REQUEST: {ctx.input or '(not provided)'}\n\n"
            f"TOOL CALLS:\n{tool_summary}"
        )

        resp = provider.judge(system=TOOL_EFFICIENCY_PROMPT, user=user_prompt)
        score = resp.get("score", 0.5)

        return CheckResult(
            name=self.name,
            score=score,
            passed=score >= self.threshold,
            threshold=self.threshold,
            evidence=resp.get("evidence", "Efficiency evaluated"),
            cost_usd=resp.get("cost_usd", 0.0),
            tokens_used=resp.get("tokens", {}),
            raw=resp,
            check_version=self.version,
        )


@register("mcp_tool_error_handling")
class MCPToolErrorHandlingCheck(BaseCheck):
    """Check that the agent handled tool errors gracefully.

    Validates that:
    - Failed tool calls were retried or an alternative was attempted
    - The agent didn't silently ignore errors
    - Error information was communicated to the user when appropriate
    """

    requires_llm = False
    version = "1.0.0"
    description = "Validate agent error recovery behavior"

    def __init__(
        self,
        require_retry: bool = True,
        max_retries: int = 3,
        threshold: float = 0.5,
    ) -> None:
        """Initialize error handling check.

        Args:
            require_retry: Whether failed calls must be retried.
            max_retries: Maximum acceptable retries for a single failure.
            threshold: Score threshold for passing.
        """
        super().__init__(threshold=threshold)
        self.require_retry = require_retry
        self.max_retries = max_retries

    def validate_context(self, ctx: Context) -> None:
        if ctx.tool_calls is None:
            raise ContextValidationError(
                "tool_calls is required for MCPToolErrorHandlingCheck"
            )

    def evaluate(self, ctx: Context) -> CheckResult:
        self.validate_context(ctx)
        assert ctx.tool_calls is not None

        failures: list[dict[str, Any]] = []
        recoveries: list[str] = []
        issues: list[str] = []

        for i, call in enumerate(ctx.tool_calls):
            if not isinstance(call, dict):
                continue
            error = call.get("error")
            if not error:
                result = call.get("result")
                if isinstance(result, dict) and result.get("success") is False:
                    error = result.get("error", "success=False")
            if not error:
                continue

            tool_name = call.get("name", call.get("tool", "unknown"))
            failures.append({"index": i, "tool": tool_name, "error": error})

            if self.require_retry:
                retried = False
                retry_count = 0
                for j in range(i + 1, len(ctx.tool_calls)):
                    next_call = ctx.tool_calls[j]
                    if isinstance(next_call, dict):
                        next_name = next_call.get(
                            "name", next_call.get("tool", "")
                        )
                        if next_name == tool_name:
                            retried = True
                            retry_count += 1

                if retried:
                    if retry_count > self.max_retries:
                        issues.append(
                            f"'{tool_name}' retried {retry_count} times "
                            f"(max {self.max_retries})"
                        )
                    else:
                        recoveries.append(
                            f"'{tool_name}' failed and was retried"
                        )
                else:
                    # Check if an alternative tool was tried
                    remaining_tools = [
                        ctx.tool_calls[j].get(
                            "name", ctx.tool_calls[j].get("tool", "")
                        )
                        for j in range(i + 1, len(ctx.tool_calls))
                        if isinstance(ctx.tool_calls[j], dict)
                    ]
                    if remaining_tools:
                        recoveries.append(
                            f"'{tool_name}' failed; agent continued "
                            f"with other tools"
                        )
                    else:
                        issues.append(
                            f"'{tool_name}' failed with no recovery attempt"
                        )

        if not failures:
            return CheckResult(
                name=self.name,
                score=1.0,
                passed=True,
                threshold=self.threshold,
                evidence="No tool errors occurred",
                check_version=self.version,
            )

        total_events = len(failures)
        score = max(
            0.0,
            (total_events - len(issues)) / total_events,
        )

        parts: list[str] = []
        parts.append(f"{len(failures)} tool error(s)")
        if recoveries:
            parts.append(f"{len(recoveries)} recovered")
        if issues:
            parts.append(f"{len(issues)} issue(s): " + "; ".join(issues))

        return CheckResult(
            name=self.name,
            score=score,
            passed=score >= self.threshold,
            threshold=self.threshold,
            evidence=". ".join(parts),
            check_version=self.version,
        )
