from __future__ import annotations

import logging
import uuid
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any

import pytest
from deepagents.backends.utils import create_file_data, file_data_to_string
from langchain_core.messages import AIMessage, AnyMessage, ToolMessage
from langsmith import testing as t
from langsmith.run_helpers import get_current_run_tree

if TYPE_CHECKING:
    from langchain_core.language_models import BaseChatModel
    from langgraph.graph.state import CompiledStateGraph

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Core trajectory data structures
# ---------------------------------------------------------------------------


@dataclass(frozen=True)
class AgentStep:
    """A step of the agent."""

    index: int
    """Start counting from 1"""
    action: AIMessage
    """AI message output from the agent. May or may not contain tool calls."""
    observations: list[ToolMessage]
    """Any observations made through tool calls."""

    def __post_init__(self) -> None:
        if self.index <= 0:
            msg = "index must be positive"
            raise ValueError(msg)


@dataclass(frozen=True)
class AgentTrajectory:
    """A trajectory of the agent."""

    steps: list[AgentStep]
    files: dict[str, str]

    @property
    def answer(self) -> str:
        """Return the text content of the last agent step."""
        return self.steps[-1].action.text

    def pretty(self) -> str:
        """Return a human-readable summary of the trajectory."""
        lines: list[str] = []
        for step in self.steps:
            lines.append(f"step {step.index}:")
            tool_calls = step.action.tool_calls
            if tool_calls:
                for tc in tool_calls:
                    name = tc.get("name")
                    args = tc.get("args")
                    lines.append(f"  - {name} {args}")
            text = step.action.text
            if text and text.strip():
                text_preview = text.strip().replace("\n", "\\n")
                lines.append(f"  text: {text_preview}")
        return "\n".join(lines)


# ---------------------------------------------------------------------------
# Assertion base classes
# ---------------------------------------------------------------------------


class SuccessAssertion:
    """Base for correctness assertions that fail the test when violated."""

    def check(self, trajectory: AgentTrajectory) -> bool:
        """Return `True` when the assertion holds.

        Args:
            trajectory: The agent trajectory to check.

        Returns:
            Whether the assertion passed.

        Raises:
            NotImplementedError: Subclasses must override this method.
        """
        raise NotImplementedError

    def describe_failure(self, trajectory: AgentTrajectory) -> str:
        """Return a human-readable explanation of why the check failed.

        Args:
            trajectory: The agent trajectory that failed the check.

        Returns:
            A description of the failure.

        Raises:
            NotImplementedError: Subclasses must override this method.
        """
        raise NotImplementedError


@dataclass(frozen=True)
class EfficiencyAssertion:
    """Base for trajectory-shape assertions that are logged but never fail."""

    def check(self, trajectory: AgentTrajectory) -> bool:
        """Return `True` when the assertion holds.

        Args:
            trajectory: The agent trajectory to check.

        Returns:
            Whether the assertion passed.

        Raises:
            NotImplementedError: Subclasses must override this method.
        """
        raise NotImplementedError

    def describe_failure(self, trajectory: AgentTrajectory) -> str:
        """Return a human-readable explanation of why the check failed.

        Args:
            trajectory: The agent trajectory that failed the check.

        Returns:
            A description of the failure.

        Raises:
            NotImplementedError: Subclasses must override this method.
        """
        raise NotImplementedError


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _strip_common_zero_width(text: str) -> str:
    """Remove common zero-width characters that can break string comparisons.

    Some models insert invisible Unicode characters (e.g., zero-width spaces) into
    strings that look like paths or URLs. This is typically a formatting heuristic
    to prevent auto-linking or to stabilize rendering in downstream UIs.

    Our eval expectations are literal substring checks. Stripping these characters
    makes the checks robust to formatting-only differences while preserving the
    semantic content of the answer.

    Args:
        text: The input string.

    Returns:
        The string with zero-width characters removed.
    """
    return text.translate(
        {
            ord("\u200b"): None,  # zero-width space
            ord("\u200c"): None,  # zero-width non-joiner
            ord("\u200d"): None,  # zero-width joiner
            ord("\ufeff"): None,  # zero-width no-break space / BOM
        }
    )


def _coerce_result_files_to_strings(raw_files: object) -> dict[str, str]:
    """Coerce the `files` value from an agent result into `dict[str, str]`.

    Args:
        raw_files: The raw files object from the agent result.

    Returns:
        A mapping of file paths to their string contents.

    Raises:
        TypeError: If the files object has an unexpected type.
    """
    if raw_files is None:
        return {}
    if not isinstance(raw_files, Mapping):
        msg = f"Expected files to be dict, got {type(raw_files)}"
        raise TypeError(msg)

    files: dict[str, str] = {}
    for path, file_data in raw_files.items():
        if not isinstance(path, str):
            msg = f"Expected file path to be str, got {type(path)}"
            raise TypeError(msg)

        if isinstance(file_data, str):
            files[path] = file_data
            continue

        if isinstance(file_data, Mapping) and "content" in file_data:
            files[path] = file_data_to_string(dict(file_data))
            continue

        msg = f"Unexpected file representation for {path}: {type(file_data)}"
        raise TypeError(msg)

    return files


# ---------------------------------------------------------------------------
# Concrete success assertions
# ---------------------------------------------------------------------------


@dataclass(frozen=True)
class FinalTextContains(SuccessAssertion):
    """Assert that the final agent text contains a given substring.

    Attributes:
        text: The substring to look for.
        case_insensitive: Whether the comparison should ignore case.
    """

    text: str
    case_insensitive: bool = False

    def check(self, trajectory: AgentTrajectory) -> bool:
        """Check that the final step text contains `self.text`.

        Args:
            trajectory: The agent trajectory to check.

        Returns:
            Whether the final text contains the expected substring.
        """
        haystack = _strip_common_zero_width(trajectory.steps[-1].action.text)
        needle = _strip_common_zero_width(self.text)
        if self.case_insensitive:
            haystack = haystack.lower()
            needle = needle.lower()
        return needle in haystack

    def describe_failure(self, trajectory: AgentTrajectory) -> str:
        """Describe why the final-text-contains check failed.

        Args:
            trajectory: The agent trajectory that failed the check.

        Returns:
            A human-readable failure description.
        """
        final_text = _strip_common_zero_width(trajectory.steps[-1].action.text)
        return f"Expected final text to contain {self.text!r} (case_insensitive={self.case_insensitive}), got: {final_text!r}"


@dataclass(frozen=True)
class FinalTextExcludes(SuccessAssertion):
    """Assert that the final agent text does NOT contain a given substring.

    Attributes:
        text: The substring that must be absent.
        case_insensitive: Whether the comparison should ignore case.
    """

    text: str
    case_insensitive: bool = False

    def check(self, trajectory: AgentTrajectory) -> bool:
        """Check that the final step text does not contain `self.text`.

        Args:
            trajectory: The agent trajectory to check.

        Returns:
            Whether the final text excludes the expected substring.
        """
        haystack = _strip_common_zero_width(trajectory.steps[-1].action.text)
        needle = _strip_common_zero_width(self.text)
        if self.case_insensitive:
            haystack = haystack.lower()
            needle = needle.lower()
        return needle not in haystack

    def describe_failure(self, trajectory: AgentTrajectory) -> str:
        """Describe why the final-text-excludes check failed.

        Args:
            trajectory: The agent trajectory that failed the check.

        Returns:
            A human-readable failure description.
        """
        final_text = _strip_common_zero_width(trajectory.steps[-1].action.text)
        return f"Expected final text NOT to contain {self.text!r} (case_insensitive={self.case_insensitive}), got: {final_text!r}"


@dataclass(frozen=True)
class FinalTextContainsAny(SuccessAssertion):
    """Assert that the final agent text contains at least ONE of the given substrings.

    Useful when a behavior can be expressed in several equivalent phrasings
    (e.g., a "value is missing" acknowledgement that could be worded as
    "unknown", "no data", "n/a", "unable to look up", "cannot be ranked",
    etc.). The any-of group still proves the behavior — a model that fakes
    or hallucinates instead of acknowledging the gap would not include any
    of these phrases.

    Attributes:
        texts: The set of substrings to look for; the check passes if any
            one of them is present.
        case_insensitive: Whether the comparison should ignore case.
    """

    texts: tuple[str, ...]
    case_insensitive: bool = False

    def check(self, trajectory: AgentTrajectory) -> bool:
        """Check that the final step text contains at least one of `self.texts`.

        Args:
            trajectory: The agent trajectory to check.

        Returns:
            Whether the final text contains any of the expected substrings.
        """
        haystack = _strip_common_zero_width(trajectory.steps[-1].action.text)
        if self.case_insensitive:
            haystack = haystack.lower()
        for text in self.texts:
            needle = _strip_common_zero_width(text)
            if self.case_insensitive:
                needle = needle.lower()
            if needle in haystack:
                return True
        return False

    def describe_failure(self, trajectory: AgentTrajectory) -> str:
        """Describe why the final-text-contains-any check failed.

        Args:
            trajectory: The agent trajectory that failed the check.

        Returns:
            A human-readable failure description.
        """
        final_text = _strip_common_zero_width(trajectory.steps[-1].action.text)
        return (
            f"Expected final text to contain at least one of "
            f"{list(self.texts)!r} (case_insensitive={self.case_insensitive}), "
            f"got: {final_text!r}"
        )


@dataclass(frozen=True)
class FinalTextMinLength(SuccessAssertion):
    """Assert that the final agent text is at least `n` chars (after strip).

    Useful for filtering out short recap-style wrap-up messages that may
    happen to contain expected substrings but aren't the substantive
    answer the user asked for.

    Attributes:
        n: Minimum length of `trajectory.steps[-1].action.text.strip()`.
    """

    n: int

    def check(self, trajectory: AgentTrajectory) -> bool:
        """Check that the stripped final text is at least `self.n` chars.

        Args:
            trajectory: The agent trajectory to check.

        Returns:
            Whether the final text meets the minimum length.
        """
        return len(trajectory.steps[-1].action.text.strip()) >= self.n

    def describe_failure(self, trajectory: AgentTrajectory) -> str:
        """Describe why the final-text-min-length check failed.

        Args:
            trajectory: The agent trajectory that failed the check.

        Returns:
            A human-readable failure description.
        """
        final_text = _strip_common_zero_width(trajectory.steps[-1].action.text)
        actual = len(final_text.strip())
        return f"Expected final text length >= {self.n}, got {actual}: {final_text!r}"


@dataclass(frozen=True)
class FileEquals(SuccessAssertion):
    """Assert that a file in the trajectory has exactly the expected content.

    Attributes:
        path: The file path to check.
        content: The expected full content of the file.
    """

    path: str
    content: str

    def check(self, trajectory: AgentTrajectory) -> bool:
        """Check that the file at `self.path` equals `self.content`.

        Args:
            trajectory: The agent trajectory to check.

        Returns:
            Whether the file content matches exactly.
        """
        return trajectory.files.get(self.path) == self.content

    def describe_failure(self, trajectory: AgentTrajectory) -> str:
        """Describe why the file-equals check failed.

        Args:
            trajectory: The agent trajectory that failed the check.

        Returns:
            A human-readable failure description.
        """
        actual = trajectory.files.get(self.path)
        if actual is None:
            return f"File {self.path!r} not found in trajectory files"
        return f"File {self.path!r} content mismatch.\nExpected:\n{self.content!r}\nActual:\n{actual!r}"


@dataclass(frozen=True)
class FileContains(SuccessAssertion):
    """Assert that a file in the trajectory contains a given substring.

    Attributes:
        path: The file path to check.
        substring: The substring to look for.
    """

    path: str
    substring: str

    def check(self, trajectory: AgentTrajectory) -> bool:
        """Check that the file at `self.path` contains `self.substring`.

        Args:
            trajectory: The agent trajectory to check.

        Returns:
            Whether the file content contains the substring.
        """
        file_content = trajectory.files.get(self.path)
        if file_content is None:
            return False
        return self.substring in file_content

    def describe_failure(self, trajectory: AgentTrajectory) -> str:
        """Describe why the file-contains check failed.

        Args:
            trajectory: The agent trajectory that failed the check.

        Returns:
            A human-readable failure description.
        """
        actual = trajectory.files.get(self.path)
        if actual is None:
            return f"File {self.path!r} not found in trajectory files"
        return (
            f"File {self.path!r} does not contain {self.substring!r}.\nActual content:\n{actual!r}"
        )


@dataclass(frozen=True)
class FileExcludes(SuccessAssertion):
    """Assert that a file in the trajectory does NOT contain a given substring.

    Attributes:
        path: The file path to check.
        substring: The substring that must be absent.
    """

    path: str
    substring: str

    def check(self, trajectory: AgentTrajectory) -> bool:
        """Check that the file at `self.path` does not contain `self.substring`.

        Args:
            trajectory: The agent trajectory to check.

        Returns:
            Whether the file content excludes the substring.
        """
        return self.substring not in trajectory.files.get(self.path, "")

    def describe_failure(self, trajectory: AgentTrajectory) -> str:
        """Describe why the file-excludes check failed.

        Args:
            trajectory: The agent trajectory that failed the check.

        Returns:
            A human-readable failure description.
        """
        actual = trajectory.files.get(self.path, "")
        return f"File {self.path!r} unexpectedly contains {self.substring!r}.\nActual content:\n{actual!r}"


@dataclass(frozen=True)
class FileAbsent(SuccessAssertion):
    """Assert that a file path does not exist in the trajectory.

    A successful deletion removes the path from state entirely (the ``files``
    channel reducer drops keys whose update value is ``None``), so a deleted
    path should not appear in ``trajectory.files`` at all. This is stricter
    than ``FileExcludes``, which only checks for an absent substring and
    therefore passes even when the file is still present.

    Attributes:
        path: The file path that must not exist.
    """

    path: str

    def check(self, trajectory: AgentTrajectory) -> bool:
        """Check that ``self.path`` is absent from the trajectory files.

        Args:
            trajectory: The agent trajectory to check.

        Returns:
            Whether the file path is absent.
        """
        return self.path not in trajectory.files

    def describe_failure(self, trajectory: AgentTrajectory) -> str:
        """Describe why the file-absent check failed.

        Args:
            trajectory: The agent trajectory that failed the check.

        Returns:
            A human-readable failure description.
        """
        actual = trajectory.files.get(self.path)
        return (
            f"Expected file {self.path!r} to be absent, but it still exists "
            f"with content:\n{actual!r}"
        )


# ---------------------------------------------------------------------------
# Concrete efficiency assertions
# ---------------------------------------------------------------------------


@dataclass(frozen=True)
class AgentSteps(EfficiencyAssertion):
    """Assert that the trajectory has exactly `n` agent steps.

    Attributes:
        n: Expected number of agent steps.
    """

    n: int

    def check(self, trajectory: AgentTrajectory) -> bool:
        """Check that the trajectory has exactly `self.n` steps.

        Args:
            trajectory: The agent trajectory to check.

        Returns:
            Whether the step count matches.
        """
        return len(trajectory.steps) == self.n

    def describe_failure(self, trajectory: AgentTrajectory) -> str:
        """Describe why the agent-steps check failed.

        Args:
            trajectory: The agent trajectory that failed the check.

        Returns:
            A human-readable failure description.
        """
        return f"Expected {self.n} agent steps, got {len(trajectory.steps)}"


@dataclass(frozen=True)
class ToolCallRequests(EfficiencyAssertion):
    """Assert that the trajectory has exactly `n` total tool call requests.

    Attributes:
        n: Expected total number of tool call requests.
    """

    n: int

    def check(self, trajectory: AgentTrajectory) -> bool:
        """Check that total tool call requests equal `self.n`.

        Args:
            trajectory: The agent trajectory to check.

        Returns:
            Whether the tool call count matches.
        """
        actual = sum(len(s.action.tool_calls) for s in trajectory.steps)
        return actual == self.n

    def describe_failure(self, trajectory: AgentTrajectory) -> str:
        """Describe why the tool-call-requests check failed.

        Args:
            trajectory: The agent trajectory that failed the check.

        Returns:
            A human-readable failure description.
        """
        actual = sum(len(s.action.tool_calls) for s in trajectory.steps)
        return f"Expected {self.n} tool call requests, got {actual}"


@dataclass(frozen=True)
class MaxToolCallRequests(EfficiencyAssertion):
    """Assert that the trajectory has AT MOST `n` total tool call requests.

    Use this when a trivial task should not invoke tools at all (or invoke
    them rarely): a model that lost its "skip for simple tasks" guidance
    and cargo-cults a planning tool on every query produces 4+ tool calls
    where 0-1 was expected.

    Attributes:
        n: Maximum allowed number of tool call requests in the trajectory.
    """

    n: int

    def check(self, trajectory: AgentTrajectory) -> bool:
        """Check that total tool call requests do not exceed `self.n`.

        Args:
            trajectory: The agent trajectory to check.

        Returns:
            Whether the tool call count is at or below the maximum.
        """
        actual = sum(len(s.action.tool_calls) for s in trajectory.steps)
        return actual <= self.n

    def describe_failure(self, trajectory: AgentTrajectory) -> str:
        """Describe why the max-tool-call-requests check failed.

        Args:
            trajectory: The agent trajectory that failed the check.

        Returns:
            A human-readable failure description.
        """
        actual = sum(len(s.action.tool_calls) for s in trajectory.steps)
        return f"Expected at most {self.n} tool call requests, got {actual}"


@dataclass(frozen=True)
class ToolCall(EfficiencyAssertion):
    """Assert that a specific tool call occurred in the trajectory.

    When `step` is `None`, all steps are searched. When `step` is given,
    only that step (1-indexed) is checked.

    Attributes:
        name: Expected tool name.
        step: Optional 1-indexed step to restrict the search to.
        args_contains: If set, the tool call args must contain these key-value pairs.
        args_equals: If set, the tool call args must equal this dict exactly.
    """

    name: str
    step: int | None = None
    args_contains: dict[str, object] | None = None
    args_equals: dict[str, object] | None = None

    def check(self, trajectory: AgentTrajectory) -> bool:
        """Check that a matching tool call exists in the trajectory.

        Args:
            trajectory: The agent trajectory to check.

        Returns:
            Whether a matching tool call was found.
        """
        return bool(self._find_matches(trajectory))

    def describe_failure(self, trajectory: AgentTrajectory) -> str:
        """Describe why the tool-call check failed.

        Args:
            trajectory: The agent trajectory that failed the check.

        Returns:
            A human-readable failure description.
        """
        step_desc = f" in step {self.step}" if self.step is not None else ""
        return f"Missing expected tool call{step_desc}: name={self.name!r}, args_contains={self.args_contains!r}, args_equals={self.args_equals!r}"

    def _matches_tool_call(self, tc: dict[str, object]) -> bool:
        """Check whether a single tool call dict matches this expectation.

        Args:
            tc: A tool call dictionary with `name` and `args` keys.

        Returns:
            Whether the tool call matches.
        """
        if tc.get("name") != self.name:
            return False
        if self.args_contains is not None:
            args = tc.get("args")
            if not isinstance(args, dict):
                return False
            if not all(args.get(k) == v for k, v in self.args_contains.items()):
                return False
        return self.args_equals is None or tc.get("args") == self.args_equals

    def _find_matches(self, trajectory: AgentTrajectory) -> list[dict[str, object]]:
        """Find tool calls matching this expectation.

        Args:
            trajectory: The agent trajectory to search.

        Returns:
            A list of matching tool call dicts.
        """
        if self.step is not None:
            if self.step > len(trajectory.steps):
                return []
            steps_to_search = [trajectory.steps[self.step - 1]]
        else:
            steps_to_search = trajectory.steps

        return [
            tc for s in steps_to_search for tc in s.action.tool_calls if self._matches_tool_call(tc)
        ]


# ---------------------------------------------------------------------------
# Factory functions (public API)
# ---------------------------------------------------------------------------


def final_text_contains(
    text: str,
    *,
    case_insensitive: bool = False,
) -> FinalTextContains:
    """Create a `FinalTextContains` success assertion.

    Args:
        text: The substring to look for in the final agent text.
        case_insensitive: Whether the comparison should ignore case.

    Returns:
        A `FinalTextContains` assertion instance.
    """
    return FinalTextContains(text=text, case_insensitive=case_insensitive)


def final_text_excludes(
    text: str,
    *,
    case_insensitive: bool = False,
) -> FinalTextExcludes:
    """Create a `FinalTextExcludes` success assertion.

    Args:
        text: The substring that must be absent from the final agent text.
        case_insensitive: Whether the comparison should ignore case.

    Returns:
        A `FinalTextExcludes` assertion instance.
    """
    return FinalTextExcludes(text=text, case_insensitive=case_insensitive)


def final_text_contains_any(
    *texts: str,
    case_insensitive: bool = False,
) -> FinalTextContainsAny:
    """Create a `FinalTextContainsAny` success assertion.

    Args:
        *texts: The substrings to look for; the check passes if any one of
            them is present.
        case_insensitive: Whether the comparison should ignore case.

    Returns:
        A `FinalTextContainsAny` assertion instance.
    """
    return FinalTextContainsAny(texts=tuple(texts), case_insensitive=case_insensitive)


def final_text_min_length(n: int) -> FinalTextMinLength:
    """Create a `FinalTextMinLength` success assertion.

    Args:
        n: Minimum length of the final agent text (after stripping).

    Returns:
        A `FinalTextMinLength` assertion instance.
    """
    return FinalTextMinLength(n=n)


def file_equals(path: str, content: str) -> FileEquals:
    """Create a `FileEquals` success assertion.

    Args:
        path: The file path to check.
        content: The expected full content of the file.

    Returns:
        A `FileEquals` assertion instance.
    """
    return FileEquals(path=path, content=content)


def file_contains(path: str, substring: str) -> FileContains:
    """Create a `FileContains` success assertion.

    Args:
        path: The file path to check.
        substring: The substring to look for.

    Returns:
        A `FileContains` assertion instance.
    """
    return FileContains(path=path, substring=substring)


def file_excludes(path: str, substring: str) -> FileExcludes:
    """Create a `FileExcludes` success assertion.

    Args:
        path: The file path to check.
        substring: The substring that must be absent.

    Returns:
        A `FileExcludes` assertion instance.
    """
    return FileExcludes(path=path, substring=substring)


def file_absent(path: str) -> FileAbsent:
    """Create a ``FileAbsent`` success assertion.

    Args:
        path: The file path that must not exist in the trajectory files.

    Returns:
        A ``FileAbsent`` assertion instance.
    """
    return FileAbsent(path=path)


def agent_steps(n: int) -> AgentSteps:
    """Create an `AgentSteps` efficiency assertion.

    Args:
        n: Expected number of agent steps.

    Returns:
        An `AgentSteps` assertion instance.
    """
    return AgentSteps(n=n)


def tool_call_requests(n: int) -> ToolCallRequests:
    """Create a `ToolCallRequests` efficiency assertion.

    Args:
        n: Expected total number of tool call requests.

    Returns:
        A `ToolCallRequests` assertion instance.
    """
    return ToolCallRequests(n=n)


def max_tool_call_requests(n: int) -> MaxToolCallRequests:
    """Create a `MaxToolCallRequests` efficiency assertion.

    Args:
        n: Maximum allowed number of tool call requests.

    Returns:
        A `MaxToolCallRequests` assertion instance.
    """
    return MaxToolCallRequests(n=n)


def tool_call(
    name: str,
    *,
    step: int | None = None,
    args_contains: dict[str, object] | None = None,
    args_equals: dict[str, object] | None = None,
) -> ToolCall:
    """Create a `ToolCall` efficiency assertion.

    Args:
        name: Expected tool name.
        step: Optional 1-indexed step to restrict the search to.
        args_contains: If set, the tool call args must contain these key-value pairs.
        args_equals: If set, the tool call args must equal this dict exactly.

    Returns:
        A `ToolCall` assertion instance.
    """
    return ToolCall(
        name=name,
        step=step,
        args_contains=args_contains,
        args_equals=args_equals,
    )


# ---------------------------------------------------------------------------
# TrajectoryScorer (two-tier builder)
# ---------------------------------------------------------------------------


@dataclass(frozen=True)
class TrajectoryScorer:
    """Two-tier assertion container for agent trajectories.

    Use `.success()` to add correctness assertions (hard-fail) and
    `.expect()` to add efficiency assertions (logged but never fail).

    Attributes:
        _success: Tuple of success assertions.
        _expectations: Tuple of efficiency assertions.
    """

    _success: tuple[SuccessAssertion, ...] = ()
    _expectations: tuple[EfficiencyAssertion, ...] = ()

    def success(self, *assertions: SuccessAssertion) -> TrajectoryScorer:
        """Append correctness assertions that hard-fail the test when violated.

        Args:
            *assertions: One or more `SuccessAssertion` instances.

        Returns:
            A new `TrajectoryScorer` with the assertions appended.
        """
        return TrajectoryScorer(
            _success=(*self._success, *assertions),
            _expectations=self._expectations,
        )

    def expect(
        self,
        *,
        agent_steps: int | None = None,
        tool_call_requests: int | None = None,
        tool_calls: list[ToolCall] | None = None,
    ) -> TrajectoryScorer:
        """Append efficiency assertions that are logged but never fail.

        Args:
            agent_steps: Expected number of agent steps.
            tool_call_requests: Expected total tool call requests.
            tool_calls: Expected tool calls with optional step pinning.

        Returns:
            A new `TrajectoryScorer` with the assertions appended.
        """
        new: list[EfficiencyAssertion] = []
        if agent_steps is not None:
            new.append(AgentSteps(n=agent_steps))
        if tool_call_requests is not None:
            new.append(ToolCallRequests(n=tool_call_requests))
        if tool_calls is not None:
            new.extend(tool_calls)
        return TrajectoryScorer(
            _success=self._success,
            _expectations=(*self._expectations, *new),
        )


# ---------------------------------------------------------------------------
# Internal: trajectory construction & assertion runner
# ---------------------------------------------------------------------------


def _trajectory_from_result(result: Mapping[str, object]) -> AgentTrajectory:
    """Build an `AgentTrajectory` from a raw agent invoke result.

    Args:
        result: The mapping returned by `agent.invoke()`.

    Returns:
        The constructed `AgentTrajectory`.

    Raises:
        TypeError: If `result['messages']` is not a list.
    """
    steps: list[AgentStep] = []
    current_step: AgentStep | None = None

    messages_obj = result.get("messages")
    if not isinstance(messages_obj, list):
        msg = f"Expected result['messages'] to be list, got {type(messages_obj)}"
        raise TypeError(msg)

    for msg_obj in messages_obj[1:]:
        if isinstance(msg_obj, AIMessage):
            if current_step is not None:
                steps.append(current_step)
            current_step = AgentStep(index=len(steps) + 1, action=msg_obj, observations=[])
        elif isinstance(msg_obj, ToolMessage):
            if current_step is not None:
                current_step.observations.append(msg_obj)

    if current_step is not None:
        steps.append(current_step)

    return AgentTrajectory(
        steps=steps,
        files=_coerce_result_files_to_strings(result.get("files")),
    )


@dataclass
class EfficiencyResult:
    """Per-test efficiency data collected during the session."""

    expected_steps: int | None
    actual_steps: int
    expected_tool_calls: int | None
    actual_tool_calls: int
    duration_s: float | None = None
    passed: bool | None = None


_on_efficiency_result: Callable[[EfficiencyResult], None] | None = None
"""Callback set by the reporter plugin to collect per-test efficiency data."""


def _log_efficiency(
    trajectory: AgentTrajectory,
    scorer: TrajectoryScorer,
) -> EfficiencyResult | None:
    """Log efficiency feedback to LangSmith and return collected data.

    Args:
        trajectory: The agent trajectory.
        scorer: The scorer containing efficiency expectations.

    Returns:
        An `EfficiencyResult` when the scorer has step or tool-call
        expectations, `None` otherwise.
    """
    actual_steps = len(trajectory.steps)
    actual_tool_calls = sum(len(s.action.tool_calls) for s in trajectory.steps)
    t.log_feedback(key="agent_steps", value=actual_steps)
    t.log_feedback(key="tool_call_requests", value=actual_tool_calls)

    expected_steps: int | None = None
    expected_tool_calls: int | None = None
    for assertion in scorer._expectations:
        if isinstance(assertion, AgentSteps):
            expected_steps = assertion.n
        elif isinstance(assertion, ToolCallRequests):
            expected_tool_calls = assertion.n

    if expected_steps is not None:
        t.log_feedback(key="expected_agent_steps", value=expected_steps)
    if expected_tool_calls is not None:
        t.log_feedback(key="expected_tool_call_requests", value=expected_tool_calls)

    if expected_steps is None and expected_tool_calls is None:
        return None

    return EfficiencyResult(
        expected_steps=expected_steps,
        actual_steps=actual_steps,
        expected_tool_calls=expected_tool_calls,
        actual_tool_calls=actual_tool_calls,
    )


def _assert_expectations(
    trajectory: AgentTrajectory,
    scorer: TrajectoryScorer,
) -> None:
    """Run all assertions in *scorer* against *trajectory*.

    Success assertions hard-fail the test via `pytest.fail`. Efficiency
    assertions are logged as feedback but never cause a test failure.

    Args:
        trajectory: The agent trajectory to validate.
        scorer: The two-tier expectation container.
    """
    eff_result = _log_efficiency(trajectory, scorer)
    if eff_result is not None and _on_efficiency_result is not None:
        _on_efficiency_result(eff_result)

    # Hard correctness checks
    success = True
    for assertion in scorer._success:
        if not assertion.check(trajectory):
            success = False
            t.log_feedback(key="correctness", value=0)
            pytest.fail(
                f"success check failed: {assertion.describe_failure(trajectory)}\n\ntrajectory:\n{trajectory.pretty()}",
                pytrace=False,
            )
    if success:
        t.log_feedback(key="correctness", value=1)


# ---------------------------------------------------------------------------
# Public entry point
# ---------------------------------------------------------------------------


def _build_invoke_inputs(
    query: str | list[AnyMessage],
    initial_files: dict[str, str] | None,
    extra_state: dict[str, Any] | None = None,
) -> dict[str, Any]:
    """Build the input payload passed to the agent invoke methods."""
    if isinstance(query, str):
        invoke_inputs: dict[str, Any] = {"messages": [{"role": "user", "content": query}]}
    else:
        invoke_inputs = {"messages": query}
    if initial_files is not None:
        invoke_inputs["files"] = {
            path: create_file_data(content) for path, content in initial_files.items()
        }
    if extra_state:
        # Shallow merge so callers can populate middleware-owned state fields
        # (e.g. `rubric` for RubricMiddleware) without clobbering `messages`/`files`.
        for key, value in extra_state.items():
            invoke_inputs[key] = value
    return invoke_inputs


def _build_logged_inputs(
    model: BaseChatModel,
    eval_metadata: dict[str, object] | None,
) -> dict[str, Any]:
    """Build the LangSmith input payload for the current test run."""
    run_tree = get_current_run_tree()
    model_str = str(getattr(model, "model", None) or getattr(model, "model_name", ""))
    logged_inputs: dict[str, Any] = {
        "test_name": run_tree.name if run_tree else "unknown",
        "model": model_str,
    }
    if eval_metadata is not None:
        logged_inputs["eval_metadata"] = eval_metadata
    return logged_inputs


def _log_run_inputs(logged_inputs: dict[str, Any]) -> None:
    """Replace LangSmith auto-captured inputs with a minimal eval payload."""
    t.log_inputs(logged_inputs)
    run_tree = get_current_run_tree()
    if run_tree is not None:
        run_tree.inputs = logged_inputs
    else:
        logger.debug(
            "run_tree is None; run_tree.inputs will not be overridden "
            "(sync_example may record auto-captured inputs)"
        )


def run_agent(
    agent: CompiledStateGraph[Any, Any],
    *,
    query: str | list[AnyMessage],
    model: BaseChatModel,
    initial_files: dict[str, str] | None = None,
    scorer: TrajectoryScorer | None = None,
    thread_id: str | None = None,
    eval_metadata: dict[str, object] | None = None,
    extra_state: dict[str, Any] | None = None,
) -> AgentTrajectory:
    """Run agent eval against the given query.

    Args:
        agent: The compiled state graph to invoke.
        query: A string prompt or list of messages.
        model: The chat model (used for logging only).
        initial_files: Optional initial files to seed the agent with.
        scorer: Optional trajectory expectations to validate.
        thread_id: Optional thread ID for the invocation.
        eval_metadata: Optional metadata to attach to the logged inputs.
        extra_state: Optional extra fields merged into the invoke input
            (e.g. `{"rubric": "..."}` for `RubricMiddleware`).

    Returns:
        The resulting `AgentTrajectory`.

    Raises:
        TypeError: If the invoke result is not a `Mapping`.
    """
    invoke_inputs = _build_invoke_inputs(query, initial_files, extra_state)

    if thread_id is None:
        thread_id = str(uuid.uuid4())
    config = {"configurable": {"thread_id": thread_id}}

    logged_inputs = _build_logged_inputs(model, eval_metadata)
    _log_run_inputs(logged_inputs)
    result = agent.invoke(invoke_inputs, config)
    t.log_outputs(result)

    if not isinstance(result, Mapping):
        msg = f"Expected invoke result to be Mapping, got {type(result)}"
        raise TypeError(msg)

    trajectory = _trajectory_from_result(result)
    if scorer is not None:
        _assert_expectations(trajectory, scorer)
    return trajectory


async def run_agent_async(
    agent: CompiledStateGraph[Any, Any],
    *,
    query: str | list[AnyMessage],
    model: BaseChatModel,
    initial_files: dict[str, str] | None = None,
    scorer: TrajectoryScorer | None = None,
    thread_id: str | None = None,
    eval_metadata: dict[str, object] | None = None,
    extra_state: dict[str, Any] | None = None,
) -> AgentTrajectory:
    """Run agent eval asynchronously against the given query.

    Args:
        agent: The compiled state graph to invoke.
        query: A string prompt or list of messages.
        model: The chat model (used for logging only).
        initial_files: Optional initial files to seed the agent with.
        scorer: Optional trajectory expectations to validate.
        thread_id: Optional thread ID for the invocation.
        eval_metadata: Optional metadata to attach to the logged inputs.
        extra_state: Optional extra fields merged into the invoke input
            (e.g. `{"rubric": "..."}` for `RubricMiddleware`).

    Returns:
        The resulting `AgentTrajectory`.

    Raises:
        TypeError: If the invoke result is not a `Mapping`.
    """
    invoke_inputs = _build_invoke_inputs(query, initial_files, extra_state)

    if thread_id is None:
        thread_id = str(uuid.uuid4())
    config = {"configurable": {"thread_id": thread_id}}

    logged_inputs = _build_logged_inputs(model, eval_metadata)
    _log_run_inputs(logged_inputs)
    result = await agent.ainvoke(invoke_inputs, config)
    t.log_outputs(result)

    if not isinstance(result, Mapping):
        msg = f"Expected invoke result to be Mapping, got {type(result)}"
        raise TypeError(msg)

    trajectory = _trajectory_from_result(result)
    if scorer is not None:
        _assert_expectations(trajectory, scorer)
    return trajectory
