"""Visibility rules for transport-only tools used inside the agent graph.

The graph/checkpoint must retain these calls because they are part of the
model's causal history. Public serializers use the projection helpers below
without mutating that history.
"""

from __future__ import annotations

from collections.abc import Iterable, Mapping
from typing import Any

INTERNAL_TOOL_NAMES = frozenset({"begin_write_file"})


def is_internal_tool_name(name: object) -> bool:
    """Return whether *name* is an exact match for a transport-only tool."""

    return isinstance(name, str) and name in INTERNAL_TOOL_NAMES


def _serialized_tool_name(tool_call: object) -> object:
    if not isinstance(tool_call, Mapping):
        return None
    name = tool_call.get("name")
    if name is not None:
        return name
    function = tool_call.get("function")
    return function.get("name") if isinstance(function, Mapping) else None


def internal_tool_call_ids(messages: Iterable[object]) -> set[str]:
    """Collect ids belonging to exact internal calls in serialized messages."""

    ids: set[str] = set()
    for message in messages:
        if not isinstance(message, Mapping):
            continue
        tool_calls = message.get("tool_calls")
        if not isinstance(tool_calls, (list, tuple)):
            continue
        for tool_call in tool_calls:
            if not is_internal_tool_name(_serialized_tool_name(tool_call)):
                continue
            call_id = tool_call.get("id") if isinstance(tool_call, Mapping) else None
            if call_id:
                ids.add(str(call_id))
    return ids


def project_outward_messages(
    messages: Iterable[Any],
    *,
    known_internal_tool_call_ids: Iterable[str] = (),
) -> list[Any]:
    """Remove transport-only call/result pairs from serialized messages."""

    materialized = list(messages)
    internal_ids = {str(call_id) for call_id in known_internal_tool_call_ids if call_id}
    projected: list[Any] = []

    for message in materialized:
        if not isinstance(message, Mapping):
            projected.append(message)
            continue

        kind = str(message.get("type") or message.get("role") or "").lower()
        if kind in {"ai", "assistant", "aimessage", "aimessagechunk"}:
            tool_calls = message.get("tool_calls")
            if isinstance(tool_calls, (list, tuple)) and tool_calls:
                internal_ids.update(internal_tool_call_ids([message]))
                visible_calls = [
                    tool_call
                    for tool_call in tool_calls
                    if not is_internal_tool_name(_serialized_tool_name(tool_call))
                ]
                if not visible_calls:
                    continue
                if len(visible_calls) != len(tool_calls):
                    message = {**message, "tool_calls": visible_calls}
        elif kind in {"tool", "toolmessage"}:
            call_id = str(message.get("tool_call_id") or "")
            if is_internal_tool_name(message.get("name")) or call_id in internal_ids:
                internal_ids.discard(call_id)
                continue

        projected.append(message)

    return projected


__all__ = [
    "INTERNAL_TOOL_NAMES",
    "internal_tool_call_ids",
    "is_internal_tool_name",
    "project_outward_messages",
]