"""provider-agnostic tool-name translation for chat-model wire boundaries.

The canonical 3tears tool name is a dotted form (``threetears.calculator``,
``3tears.admin.agent_management``). Several upstream model APIs validate
tool names against ``^[a-zA-Z0-9_-]{1,128}$`` and reject the dot:

- Anthropic Messages API (``claude-*-direct``): ``invalid_request_error``
  with the regex above explicitly in the message.
- AWS Bedrock-routed Anthropic + most OpenAI-compat backends accessed via
  OpenRouter: same regex, surfaced as a 4xx from the upstream provider.
- OpenAI Responses / chat-completions tools surface: same regex, sometimes
  different text but the constraint is identical.

This module provides the wire-side translation primitives so EVERY chat
provider can wear the same translation layer without duplicating it. The
two pieces are:

- :class:`NameMangledToolProxy` -- a LangChain :class:`BaseTool` whose
  ``.name`` is the dot-replaced wire form, delegating execution to the
  original tool. The proxy preserves description and ``args_schema``
  verbatim, so the LLM still sees the unchanged tool spec aside from
  the name.
- :func:`build_name_translation` -- builds a ``(wire_tools, reverse_map)``
  tuple from an application-side tool list. ``wire_tools`` is the list
  the chat model binds to; ``reverse_map`` maps each underscored wire
  name back to its canonical dotted form so streaming / non-streaming
  responses can be un-translated before reaching application code.
- :func:`reverse_translate_message` -- mutates an
  ``AIMessage`` / ``AIMessageChunk`` in place, rewriting every
  ``tool_calls`` / ``tool_call_chunks`` / ``invalid_tool_calls`` name
  field through ``reverse_map``. Provider-specific
  ``_astream`` / ``_agenerate`` wrappers call this so consumers see
  canonical dotted tool names regardless of which provider's
  validator forced the wire-side rename.

Per-provider integration is a thin subclass that:

1. Overrides ``bind_tools`` to call :func:`build_name_translation` and
   store the resulting reverse map on ``_name_reverse_map`` (a
   :class:`PrivateAttr` on the subclass).
2. Overrides the response surface to call
   :func:`reverse_translate_message` on every message it returns. This
   must cover BOTH the streaming and non-streaming public entry points,
   because LangChain's ``ainvoke`` / ``invoke`` route internally through
   the PROTECTED ``_astream`` / ``_stream`` (not ``_agenerate`` /
   ``_generate``) whenever ``_should_stream()`` is true — e.g. when a
   streaming callback is attached (LangGraph's ``astream_events`` tap).
   The concrete subclasses therefore override the public ``astream`` +
   ``ainvoke`` + ``invoke`` (post-processing the result) AND
   ``_agenerate`` (for the pure non-streaming path). Overriding the
   protected ``_astream`` directly is deliberately avoided — wrapping it
   in another async generator drops ``on_chat_model_stream`` callbacks
   (see the ``astream`` docstrings in the provider modules).

See :mod:`threetears.models.providers.openrouter` and
:mod:`threetears.models.providers.anthropic` for the two concrete
integrations.

``config`` propagation bug (found live: every REAL 3tears builtin tool -- e.g.
``threetears.calculator`` -- raised ``TypeError: StructuredTool._arun() missing 1 required
keyword-only argument: 'config'`` whenever actually invoked through a proxied (dotted) name,
across every provider that uses this module, not just one). Root cause: ``BaseTool.arun()``/
``run()`` only forward ``config`` (and ``run_manager``) into ``self._arun(...)``/``self._run(...)``
when THAT method's own signature declares a ``RunnableConfig``-typed parameter -- it introspects
``self._arun`` via type hints, not just "does a config exist somewhere". ``NameMangledToolProxy``'s
``_arun``/``_run`` never declared one, so LangChain's own machinery never gave the proxy a
``config`` to forward, and the proxy's body called ``self._delegate._arun(*args, **kwargs)``
directly -- bypassing the delegate's own public ``arun()``/``run()`` wrapper (deliberately, per the
double-callback rationale in the method docstrings below) with no ``config`` in ``kwargs`` at all.
That's harmless for a delegate whose ``_arun`` doesn't require it (a plain custom ``BaseTool``
subclass), but every REAL 3tears builtin is a ``StructuredTool`` (built by
:func:`~threetears.agent.tools.langchain_adapter.to_langchain_tool`), and LangChain's own
``StructuredTool._arun`` declares ``config: RunnableConfig`` as a REQUIRED keyword-only parameter
-- so every single proxied builtin tool call failed. Fixed by declaring ``config``/``run_manager``
on the proxy's own ``_arun``/``_run`` (so the caller's introspection finds and supplies them), then
forwarding them to the delegate ONLY if the delegate's own method actually wants them -- the exact
same conditional introspection ``BaseTool.arun()``/``run()`` themselves use, applied one level
deeper.
"""

from __future__ import annotations

import inspect
from typing import Any, get_type_hints

from langchain_core.callbacks import AsyncCallbackManagerForToolRun, CallbackManagerForToolRun
from langchain_core.runnables import RunnableConfig
from langchain_core.tools import BaseTool
from pydantic import PrivateAttr

__all__ = [
    "NameMangledToolProxy",
    "build_name_translation",
    "forward_translate_input",
    "forward_translate_message",
    "mangle_tool_name",
    "reverse_translate_message",
]


def mangle_tool_name(name: str) -> str:
    """translate a dotted canonical tool name to its underscored wire form.

    Applied before sending tool specs to providers whose validators
    reject dots. The transform is the simplest one that satisfies
    every observed validator (``^[a-zA-Z0-9_-]{1,128}$``); existing
    underscores in the name are preserved so the round-trip is
    lossless when paired with the reverse map built by
    :func:`build_name_translation`.

    :param name: canonical tool name (may contain dots)
    :ptype name: str
    :return: dot-replaced form suitable for strict provider regex
        validators (Bedrock, Anthropic-direct, OpenAI Responses,
        OpenRouter)
    :rtype: str
    """
    return name.replace(".", "_")


def _delegate_context_kwargs(
    delegate_method: Any,
    config: RunnableConfig,
    run_manager: Any,
) -> dict[str, Any]:
    """Build the subset of ``{<config param name>: config, "run_manager": run_manager}``
    ``delegate_method`` actually declares -- mirrors ``BaseTool.arun``/``run``'s own
    introspection (a ``RunnableConfig``-type-hinted parameter for config, bare presence for
    ``run_manager``) one level deeper, so a delegate gets exactly what its OWN signature wants,
    never an unexpected kwarg and never a silently-missing required one.
    """
    forwarded: dict[str, Any] = {}
    try:
        hints = get_type_hints(delegate_method)
    except Exception:  # prawduct:allow prawduct/broad-except -- signature introspection on an
        # arbitrary delegate must never break dispatch; worst case config isn't forwarded and
        # the delegate raises its own clear error, same as before this fix existed.
        hints = {}
    for name, type_ in hints.items():
        if type_ is RunnableConfig:
            forwarded[name] = config
            break
    if "run_manager" in inspect.signature(delegate_method).parameters:
        forwarded["run_manager"] = run_manager
    return forwarded


class NameMangledToolProxy(BaseTool):
    """LangChain :class:`BaseTool` whose ``.name`` is the wire form,
    delegating execution to a dotted-named original tool.

    Used by every provider's name-translating chat subclass to swap
    each application-side tool for a wire-side proxy at
    ``bind_tools`` time. The proxy keeps the ``description`` and
    ``args_schema`` of the original (those flow to the LLM
    unchanged) but exposes the dot-replaced name so the bound tool
    spec passes strict provider validators.

    Application code never receives this proxy directly -- it lives
    only inside the bound runnable's tool list, and the response
    un-translation layer (:func:`reverse_translate_message`) rewrites
    tool-call names back to the dotted form before any consumer code
    sees them.
    """

    name: str
    description: str
    args_schema: type[Any] | None = None
    _delegate: BaseTool = PrivateAttr()

    def __init__(self, *, delegate: BaseTool, mangled_name: str) -> None:
        """proxy initializer that copies description/args from the delegate.

        :param delegate: original tool whose execution to forward to
        :ptype delegate: BaseTool
        :param mangled_name: dot-replaced wire name to expose
        :ptype mangled_name: str
        """
        super().__init__(
            name=mangled_name,
            description=delegate.description,
            args_schema=delegate.args_schema,
        )
        self._delegate = delegate

    @property
    def canonical_name(self) -> str:
        """the delegate's un-mangled, dotted canonical name.

        Public accessor so callers who substitute proxies ahead of a chat model's own
        ``bind_tools`` (e.g. :mod:`threetears.models.providers._claude_cli`) can report the
        canonical name to their own observability/tracking, without reaching into ``_delegate``.

        :return: canonical dotted tool name
        :rtype: str
        """
        return self._delegate.name

    async def _arun(
        self,
        *args: Any,
        config: RunnableConfig,
        run_manager: AsyncCallbackManagerForToolRun | None = None,
        **kwargs: Any,
    ) -> Any:
        """forward async execution to the delegate's ``_arun``.

        Calling ``delegate._arun`` rather than ``delegate.arun`` is
        deliberate: ``arun`` is the LangChain wrapper that runs
        callbacks/validation around the subclass's ``_arun`` body.
        Our proxy already wears that wrapper at the proxy level
        (LangChain calls ``proxy.arun`` -> ``proxy._arun``); calling
        ``delegate.arun`` here would double-fire callbacks and
        re-process the input.

        ``config``/``run_manager`` are declared explicitly (rather than absorbed into
        ``**kwargs``) so ``BaseTool.arun``'s own introspection of THIS method finds a
        ``RunnableConfig``-typed parameter and actually supplies them -- see the "config
        propagation bug" note in this module's docstring. Forwarded to the delegate's ``_arun``
        ONLY if ITS OWN signature wants them (a real 3tears builtin is a ``StructuredTool``,
        whose ``_arun`` requires ``config``; a plain custom ``BaseTool`` may not accept it at
        all, so forwarding unconditionally would break that case with an unexpected-kwarg error).

        :param args: positional forwarded to the delegate
        :ptype args: Any
        :param config: this call's ``RunnableConfig``, forwarded only if the delegate wants it
        :ptype config: RunnableConfig
        :param run_manager: this call's callback manager, forwarded only if the delegate wants it
        :ptype run_manager: AsyncCallbackManagerForToolRun | None
        :param kwargs: keyword forwarded to the delegate
        :ptype kwargs: Any
        :return: delegate result
        :rtype: Any
        """
        kwargs.update(_delegate_context_kwargs(self._delegate._arun, config, run_manager))
        return await self._delegate._arun(*args, **kwargs)

    def _run(
        self,
        *args: Any,
        config: RunnableConfig,
        run_manager: CallbackManagerForToolRun | None = None,
        **kwargs: Any,
    ) -> Any:
        """forward sync execution to the delegate's ``_run``.

        Same double-wrapper rationale, and same conditional ``config``/``run_manager``
        forwarding, as :meth:`_arun`.

        :param args: positional forwarded to the delegate
        :ptype args: Any
        :param config: this call's ``RunnableConfig``, forwarded only if the delegate wants it
        :ptype config: RunnableConfig
        :param run_manager: this call's callback manager, forwarded only if the delegate wants it
        :ptype run_manager: CallbackManagerForToolRun | None
        :param kwargs: keyword forwarded to the delegate
        :ptype kwargs: Any
        :return: delegate result
        :rtype: Any
        """
        kwargs.update(_delegate_context_kwargs(self._delegate._run, config, run_manager))
        return self._delegate._run(*args, **kwargs)


def build_name_translation(
    tools: list[Any],
) -> tuple[list[Any], dict[str, str]]:
    """build the wire-side translated tool list + reverse map for un-translation.

    Accepts two input shapes that both occur in production:

    - **BaseTool list** (3tears agent runtime path). Each
      :class:`BaseTool` whose ``.name`` contains a dot is wrapped
      in a :class:`NameMangledToolProxy` carrying the underscored
      wire name; the original tool is kept as the proxy's
      delegate so execution still hits the unchanged
      ``_run`` / ``_arun``.
    - **Tool-spec dict list** (gateway path). Each dict is in the
      provider-native shape that crosses the NATS gateway boundary
      (already serialized for the upstream API). The dict's
      ``name`` field is the mangle target. Two layouts are
      handled: the OpenAI ``{"type":"function","function":{"name":...}}``
      shape and the Anthropic / 3tears-canonical ``{"name":...}``
      flat shape. A shallow-copied dict is returned so the caller's
      original list is not mutated.

    Tools whose canonical name has no dot pass through unchanged
    (no translation needed). The reverse map keys on the
    underscored form so per-provider chat subclasses can rewrite
    ``tool_call`` names in streaming and non-streaming responses.

    :param tools: application-side bind_tools input
        (BaseTool objects OR provider-native tool spec dicts)
    :ptype tools: list[Any]
    :return: tuple of (wire-side tool list, underscored -> dotted
        reverse map)
    :rtype: tuple[list[Any], dict[str, str]]
    """
    wire_tools: list[Any] = []
    reverse_map: dict[str, str] = {}
    for tool in tools:
        if isinstance(tool, dict):
            wire_tool, mapping = _translate_dict_tool(tool)
            wire_tools.append(wire_tool)
            reverse_map.update(mapping)
            continue
        # BaseTool path
        canonical = tool.name
        if "." not in canonical:
            wire_tools.append(tool)
            continue
        mangled = mangle_tool_name(canonical)
        wire_tools.append(
            NameMangledToolProxy(
                delegate=tool,
                mangled_name=mangled,
            )
        )
        reverse_map[mangled] = canonical
    return wire_tools, reverse_map


def _translate_dict_tool(
    tool: dict[str, Any],
) -> tuple[dict[str, Any], dict[str, str]]:
    """translate a single provider-native tool-spec dict.

    Two dict shapes appear in practice:

    - OpenAI / OpenAI-compat chat-completions ``tools`` entry:
      ``{"type": "function", "function": {"name": "...", "description": "...", "parameters": {...}}}``
      -- the name lives at ``["function"]["name"]``.
    - Anthropic Messages and 3tears-canonical:
      ``{"name": "...", "description": "...", "input_schema": {...}}``
      -- the name lives at the top level.

    Other layouts (no recognised name field) pass through unchanged
    so a future provider extension does not silently drop tools.

    :param tool: provider-native tool spec
    :ptype tool: dict[str, Any]
    :return: tuple of (translated tool spec, single-entry reverse
        map for this tool; empty when no translation was needed)
    :rtype: tuple[dict[str, Any], dict[str, str]]
    """
    # detect the OpenAI-shape nesting first; otherwise fall back to
    # flat ``name`` (Anthropic / canonical).
    fn_block = tool.get("function") if isinstance(tool.get("function"), dict) else None
    if fn_block is not None:
        canonical = fn_block.get("name")
        if not isinstance(canonical, str) or "." not in canonical:
            return tool, {}
        mangled = mangle_tool_name(canonical)
        new_fn = {**fn_block, "name": mangled}
        new_tool = {**tool, "function": new_fn}
        return new_tool, {mangled: canonical}
    canonical = tool.get("name")
    if not isinstance(canonical, str) or "." not in canonical:
        return tool, {}
    mangled = mangle_tool_name(canonical)
    new_tool = {**tool, "name": mangled}
    return new_tool, {mangled: canonical}


def reverse_translate_message(
    message: Any,
    reverse_map: dict[str, str],
) -> None:
    """rewrite tool-call names on ``message`` from wire to canonical form.

    Mutates in place. Called for every chunk yielded from a chat
    model's ``_astream`` and every message in ``_agenerate``'s
    result. Touches three name-bearing fields:

    - ``tool_call_chunks`` -- partial streamed tool calls; the
      ``name`` field arrives once at the start of each call,
      subsequent chunks carry only ``args``.
    - ``tool_calls`` -- the merged, well-formed tool calls
      consumers iterate.
    - ``invalid_tool_calls`` -- the recovery target for malformed
      streaming; consumers attempt to
      re-parse them.

    No-op when the message has no tool-call fields or when the
    reverse map is empty (no tools were bound, so nothing to
    translate).

    :param message: target chat message (``AIMessage`` or
        ``AIMessageChunk``); duck-typed via attribute access
    :ptype message: Any
    :param reverse_map: underscored -> dotted reverse map produced
        by :func:`build_name_translation`
    :ptype reverse_map: dict[str, str]
    """
    if not reverse_map:
        return
    for tc_chunk in getattr(message, "tool_call_chunks", None) or []:
        name = tc_chunk.get("name")
        if name and name in reverse_map:
            tc_chunk["name"] = reverse_map[name]
    for tc in getattr(message, "tool_calls", None) or []:
        name = tc.get("name")
        if name and name in reverse_map:
            tc["name"] = reverse_map[name]
    for tc in getattr(message, "invalid_tool_calls", None) or []:
        name = tc.get("name")
        if name and name in reverse_map:
            tc["name"] = reverse_map[name]


def forward_translate_message(message: Any) -> Any:
    """return a copy of ``message`` with dotted tool-call names mangled to wire form.

    Forward mirror of :func:`reverse_translate_message` for the
    OUTBOUND direction. The reverse pass rewrites a provider RESPONSE
    (wire -> canonical) and may mutate in place -- the response is
    freshly minted by the provider, so there is no application state to
    corrupt. The forward pass rewrites messages about to be SENT -- the
    conversation history, which IS application state -- so mutating in
    place would leak the wire form back into the caller's message list
    and break dispatch / logging / persistence that key on the
    canonical dotted name. This therefore returns a shallow COPY
    whenever a rename is needed and the original ``message`` otherwise.

    A dotted tool-call name reaches the wire in two ways
    :func:`build_name_translation` (which mangles only outgoing tool
    SPECS) does not cover: an ``AIMessage`` carrying a prior round's
    ``tool_calls`` re-sent on the next round, and a model hallucination
    or dotted MCP tool that never matched a bound spec. Both fail every
    provider's ``^[a-zA-Z0-9_-]`` tool-name validator. Every dotted name
    in ``tool_call_chunks`` / ``tool_calls`` / ``invalid_tool_calls`` is
    mangled via :func:`mangle_tool_name` -- the same transform
    :func:`build_name_translation` applies to specs -- so a bound tool's
    history entry round-trips losslessly through the per-bind reverse
    map.

    :param message: outbound chat message (``AIMessage`` /
        ``AIMessageChunk``); duck-typed via attribute access
    :ptype message: Any
    :return: a shallow copy with mangled tool-call names when the
        message carries a dotted name, otherwise ``message`` unchanged
    :rtype: Any
    """
    updates: dict[str, list[Any]] = {}
    for field in ("tool_call_chunks", "tool_calls", "invalid_tool_calls"):
        entries = getattr(message, field, None)
        if not entries:
            continue
        new_entries: list[Any] = []
        field_changed = False
        for entry in entries:
            name = entry.get("name") if isinstance(entry, dict) else None
            if isinstance(name, str) and "." in name:
                new_entries.append({**entry, "name": mangle_tool_name(name)})
                field_changed = True
            else:
                new_entries.append(entry)
        if field_changed:
            updates[field] = new_entries
    if not updates:
        return message
    return message.model_copy(update=updates)


def forward_translate_input(payload: Any) -> Any:
    """mangle dotted tool-call names on an outbound chat-model input.

    Applies :func:`forward_translate_message` to every message in a
    message-sequence ``payload`` (the shape the tool loop and the
    gateway send to a chat model), returning a NEW list only when at
    least one message needed a rename so the caller's list is never
    mutated. A non-``list`` ``payload`` -- a bare prompt string or a
    ``PromptValue``, which carries no tool-call history -- is returned
    unchanged.

    Applied at the four public / protected entry points a provider
    wrapper overrides (``astream`` / ``ainvoke`` / ``invoke`` /
    ``_agenerate``) before delegating to ``super()``. Idempotent: a
    second pass over already-mangled (dotless) names is a no-op that
    returns the same object, so the overlapping internal routings
    between those entry points cannot double-translate.

    :param payload: chat-model input (``list`` of messages, a string,
        or a ``PromptValue``)
    :ptype payload: Any
    :return: the payload with dotted tool-call names mangled to wire
        form (a fresh list when a rename occurred), else ``payload``
        unchanged
    :rtype: Any
    """
    if not isinstance(payload, list):
        return payload
    translated: list[Any] = []
    changed = False
    for message in payload:
        new_message = forward_translate_message(message)
        if new_message is not message:
            changed = True
        translated.append(new_message)
    return translated if changed else payload
