"""Keyword-based local checks."""

from __future__ import annotations

from agentbyte.eval.judges.reference import ExactMatchJudge
from agentbyte.eval.types import AnswerStrategy, EvalTrajectory

from .types import CheckResult, EvalCheck


def keyword_check(
    *keywords: str,
    case_sensitive: bool = False,
    source_filter: str | None = None,
    answer_strategy: AnswerStrategy = "last_assistant",
) -> EvalCheck:
    """Return a check that verifies all keywords appear in the extracted response.

    Defaults to ``"last_assistant"`` because keyword checks are always asking
    "did the agent say X?" — semantically an AssistantMessage concern. Pass a
    different ``answer_strategy`` only when checking non-assistant content.
    Pass ``source_filter`` to target a specific agent in multi-agent trajectories
    (e.g. ``source_filter="writer"`` in a Researcher→Writer→Reviewer pipeline).
    """
    extractor = ExactMatchJudge(answer_strategy=answer_strategy, source_filter=source_filter)

    async def _check(trajectory: EvalTrajectory) -> CheckResult:
        response = extractor.extract_answer(trajectory)
        haystack = response if case_sensitive else response.lower()
        expected = list(keywords if case_sensitive else [word.lower() for word in keywords])
        missing = [word for word in expected if word not in haystack]
        if missing:
            return CheckResult(
                passed=False,
                reason=f"Missing keywords: {missing}",
                check_name="keyword_check",
            )
        return CheckResult(
            passed=True,
            reason="All keywords present",
            check_name="keyword_check",
        )

    return _check


__all__ = ["keyword_check"]
