"""Tool-call checks for evaluation trajectories."""

from __future__ import annotations

from agentbyte.messages import AssistantMessage

from agentbyte.eval.types import EvalTrajectory, ExpectedToolCall

from .types import CheckResult, EvalCheck


async def tool_calls_present(trajectory: EvalTrajectory) -> CheckResult:
    """Verify that all expected tools were requested at least once."""
    expected = trajectory.task.expected_tool_calls or []
    if not expected:
        return CheckResult(
            passed=True,
            reason="No expected tool calls declared",
            check_name="tool_calls_present",
        )

    actual_names = {
        tool_call.tool_name
        for message in trajectory.messages
        if isinstance(message, AssistantMessage) and message.tool_calls
        for tool_call in message.tool_calls
    }
    missing = [item.name for item in expected if item.name not in actual_names]
    if missing:
        return CheckResult(
            passed=False,
            reason=f"Missing expected tool calls: {missing}",
            check_name="tool_calls_present",
        )
    return CheckResult(
        passed=True,
        reason="All expected tool calls were requested",
        check_name="tool_calls_present",
    )


def tool_calls_present_check(*tool_names: str) -> EvalCheck:
    """Build a tool-presence check for explicit tool names."""

    async def _check(trajectory: EvalTrajectory) -> CheckResult:
        expected = [ExpectedToolCall(name=name) for name in tool_names]
        original = trajectory.task.expected_tool_calls
        patched_task = trajectory.task.model_copy(
            update={"expected_tool_calls": expected}
        )
        patched = trajectory.model_copy(update={"task": patched_task})
        result = await tool_calls_present(patched)
        restored_reason = result.reason
        if original is not None and not tool_names:
            restored_reason = result.reason
        return result.model_copy(update={"reason": restored_reason})

    return _check


async def tool_call_args_match(trajectory: EvalTrajectory) -> CheckResult:
    """Verify expected tool call arguments against assistant tool requests."""
    expected = trajectory.task.expected_tool_calls or []
    if not expected:
        return CheckResult(
            passed=True,
            reason="No expected tool calls declared",
            check_name="tool_call_args_match",
        )

    actual_calls: list[tuple[str, dict[str, object]]] = []
    for message in trajectory.messages:
        if not isinstance(message, AssistantMessage) or not message.tool_calls:
            continue
        for tool_call in message.tool_calls:
            actual_calls.append((tool_call.tool_name, dict(tool_call.parameters)))

    missing_or_mismatched: list[str] = []
    for item in expected:
        if item.arguments is None:
            if not any(name == item.name for name, _ in actual_calls):
                missing_or_mismatched.append(f"{item.name} (missing)")
            continue
        matched = False
        for name, arguments in actual_calls:
            if name != item.name:
                continue
            if all(arguments.get(key) == value for key, value in item.arguments.items()):
                matched = True
                break
        if not matched:
            missing_or_mismatched.append(f"{item.name} args {item.arguments}")

    if missing_or_mismatched:
        return CheckResult(
            passed=False,
            reason=f"Missing or mismatched tool args: {missing_or_mismatched}",
            check_name="tool_call_args_match",
        )
    return CheckResult(
        passed=True,
        reason="All expected tool call arguments matched",
        check_name="tool_call_args_match",
    )


def tool_call_args_match_check(expected: list[ExpectedToolCall]) -> EvalCheck:
    """Build a tool-argument check for explicit expected tool calls."""

    async def _check(trajectory: EvalTrajectory) -> CheckResult:
        patched_task = trajectory.task.model_copy(
            update={"expected_tool_calls": expected}
        )
        patched = trajectory.model_copy(update={"task": patched_task})
        return await tool_call_args_match(patched)

    return _check


__all__ = [
    "tool_calls_present",
    "tool_calls_present_check",
    "tool_call_args_match",
    "tool_call_args_match_check",
]
