"""Live-agent tool-list maintenance after MCP (re)discovery: refreshing an
AIAgent's tools/tool names, preserving the cached tools[] prefix across rebuilds,
and re-injecting post-build tools."""

import logging
import json
import threading
from typing import Optional
from tools.mcp_tool_common import _core

logger = logging.getLogger("tools.mcp_tool")

# Serializes in-place swaps of ``agent.tools`` / ``agent.valid_tool_names`` by
# the reload RPC, gateway reload and late-binding refresh thread; the run loop
# reads them during tool iteration and must never see a half-updated pair.
_agent_tools_lock = threading.Lock()


def _def_name(tool_def: dict) -> str:
    return (tool_def.get("function") or {}).get("name", "")


def _agent_tool_defs(agent) -> list:
    return list(getattr(agent, "tools", None) or [])


def agent_tool_names(agent) -> list:
    """Names of ``agent.tools`` in wire order (unnamed entries skipped)."""
    return [name for name in map(_def_name, _agent_tool_defs(agent)) if name]


def _resolve_refresh_toolsets(agent, enabled_override, disabled_override):
    """Explicit reloads pass freshly-resolved toolsets (so a server just ENABLED in config is
    picked up) and the agent's selection is updated to match; automatic paths pass nothing
    and reuse the build-time selection."""
    enabled = getattr(agent, "enabled_toolsets", None)
    disabled = getattr(agent, "disabled_toolsets", None)
    if enabled_override is not None or disabled_override is not None:
        enabled = enabled_override if enabled_override is not None else enabled
        if disabled_override is not None:
            from toolsets import session_disabled_toolsets
            disabled = session_disabled_toolsets(disabled_override, getattr(agent, "platform", None))
        agent.enabled_toolsets, agent.disabled_toolsets = enabled, disabled
    return enabled, disabled


def _tool_defs_content_changed(agent, new_defs: list) -> bool:
    """Byte-level diff of the serialized tool arrays (dynamic schemas change CONTENT under
    stable names); False if either side fails to serialize."""
    try:
        dump = lambda defs: json.dumps(defs, sort_keys=True, separators=(",", ":"), default=str)
        return dump(_agent_tool_defs(agent)) != dump(new_defs)
    except Exception:
        return False


def _drop_session_tools(agent, new_defs: list, new_names: set) -> tuple:
    from toolsets import agent_tool_drops
    drops = agent_tool_drops(agent)
    if not drops:
        return new_defs, new_names
    return [entry for entry in new_defs if _def_name(entry) not in drops], new_names - drops


def _publish_tool_snapshot(
    agent, new_defs: list, new_names: set, *, snapshot_generation: int,
    staged_engine_names: set, content_aware: bool, prefix_registered: Optional[set]) -> Optional[set]:
    """Single atomic read-diff-publish under ``_agent_tools_lock`` so ``added`` matches what
    was published and a stale (older-generation) rebuild can't overwrite a newer one. Returns
    the added names, or None when nothing was published (unchanged, or a newer snapshot won)."""
    with _agent_tools_lock:
        # Tolerate an agent that never set the generation (or a non-int mock).
        published_gen = getattr(agent, "_tool_snapshot_generation", -1)
        published_gen = published_gen if isinstance(published_gen, int) else -1
        if snapshot_generation < published_gen:
            return None  # a newer snapshot already won
        current_defs = _agent_tool_defs(agent)
        current = {_def_name(t) for t in current_defs}
        if prefix_registered is not None:
            new_defs, new_names = _merge_preserving_prefix(current_defs, new_defs, prefix_registered)
        new_defs, new_names = _drop_session_tools(agent, new_defs, new_names)
        # Record the generation even when unchanged so an in-flight older caller can't clobber.
        agent._tool_snapshot_generation = max(published_gen, snapshot_generation)
        # Same name set: nothing to publish, unless the caller is content_aware
        # (explicit reload, compaction) and the serialized schema bytes moved.
        if new_names == current and not (content_aware and _tool_defs_content_changed(agent, new_defs)):
            return None
        agent.tools = new_defs
        agent.valid_tool_names = new_names
        # Publish context-engine routing names atomically with the snapshot.
        engine_names = getattr(agent, "_context_engine_tool_names", None)
        if isinstance(engine_names, set):
            engine_names.clear()
            engine_names.update(staged_engine_names)
        return new_names - current


def refresh_agent_mcp_tools(
    agent, *, enabled_override=None, disabled_override=None, quiet_mode: bool = True,
    content_aware: bool = False, preserve_prefix: bool = False) -> set:
    """Re-derive an already-built agent's tool snapshot from the live registry; returns the
    newly-added tool names (empty when unchanged). The agent snapshots ``agent.tools`` at build
    time, so servers that connect later (slow OAuth, ``/reload-mcp``) are invisible until
    rebuilt. Shared by the TUI RPC, gateway reload, late-binding thread and between-turns
    refresh: respects the toolset filter, diffs by tool NAME (a count compare misses an
    equal-size swap) and, for ``content_aware`` callers, by serialized schema bytes, re-injects the memory-provider / context-engine tools ``agent_init``
    appends after ``get_tool_definitions`` plus guarded session capabilities, and publishes
    ``(tools, valid_tool_names)`` together.

    ``preserve_prefix``: for rebuilds inside a live conversation the tool array is a cached
    request prefix and any moved byte re-prefills the whole history — existing tools keep their
    slot and the schema bytes already sent (a refreshed description is rendered ahead of the
    conversation), a still-registered tool whose ``check_fn`` merely flapped is carried forward
    (``check_fn`` gates exposure, never invocation), a deregistered tool is dropped, new tools
    append at the tail. The caller owns the prompt-cache contract.

    With ``preserve_prefix`` off (explicit ``/reload-mcp``, and a rebuild before the
    first turn) a same-name schema change is published. The caller has already
    accepted the cache break, or nothing has been cached yet."""
    from model_tools import get_tool_definitions
    from tools.registry import registry
    enabled, disabled = _resolve_refresh_toolsets(agent, enabled_override, disabled_override)
    # Generation captured BEFORE the slow get_tool_definitions call (a slower caller holding an
    # OLDER set must not clobber a newer one); definitions computed OUTSIDE the lock.
    snapshot_generation = registry._generation
    new_defs = list(get_tool_definitions(enabled_toolsets=enabled, disabled_toolsets=disabled, quiet_mode=quiet_mode) or [])
    new_names = {_def_name(t) for t in new_defs}
    # Post-build families re-appended on LOCALS only; live attributes untouched until publish.
    staged_engine_names = _reinject_post_build_tools(agent, new_defs, new_names)
    _reinject_authorized_dynamic_tools(agent, new_defs, new_names)
    # Registry membership is read OUTSIDE ``_agent_tools_lock``: taking ``registry._lock``
    # under the tools lock would be the first nesting of the two.
    prefix_registered: Optional[set] = None
    if preserve_prefix:
        try:
            prefix_registered = {entry.name for entry in registry.get_all_entries()}
        except Exception:
            pass  # fail open to the plain rebuild
    else:
        # Explicit reload and pre-first-turn rebuild. A same-name description
        # change has to publish; between-turns refresh keeps preserve_prefix.
        content_aware = True
    added = _publish_tool_snapshot(
        agent, new_defs, new_names, snapshot_generation=snapshot_generation,
        staged_engine_names=staged_engine_names, content_aware=content_aware, prefix_registered=prefix_registered)
    if added is None:
        return set()
    persist_agent_tool_names(agent)  # re-pin so a rebuild after agent-cache eviction restores this order
    return added


def reprobe_tool_availability() -> None:
    """Explicit ``/reload-mcp`` hatch out of the tools[] freeze: drop the ``check_fn`` verdict
    cache AND the ``get_tool_definitions`` memo (keyed on registry generation, so it would
    otherwise replay the stale verdicts)."""
    from model_tools import _clear_tool_defs_cache
    from tools.registry import invalidate_check_fn_cache
    invalidate_check_fn_cache()
    _clear_tool_defs_cache()


def tool_pin_version() -> str:
    """The code identity a tools[] pin was built by (checkout/build sha, else the release version).
    Cached per process: an updated checkout only reaches a process through a restart."""
    from hermes_cli import __version__
    from hermes_cli.build_info import get_code_identity
    identity = get_code_identity()
    return identity.get("sha") or identity.get("version") or __version__


def persist_agent_tool_names(agent) -> None:
    """Best-effort: write ``agent.tools`` to the session row (freeze pin). The full definitions,
    keyed by the code that built them: another process or surface derives different bytes."""
    db = getattr(agent, "_session_db", None)
    session_id = getattr(agent, "session_id", None)
    if not db or not session_id:
        return
    try:
        db.update_session_tool_names(session_id, {"version": tool_pin_version(), "tools": _agent_tool_defs(agent)})
    except Exception:
        logger.debug("tool_names persist skipped", exc_info=True)


def _config_permitted_names(agent) -> set:
    """Tool names this agent's toolset selection allows before ``check_fn``: all a pin may carry
    forward. A client-surface toolset counts as allowed (only its client can add it, so its absence
    here is no config choice); ``disabled_toolsets`` still strips it."""
    from model_tools import _select_tool_names
    from toolsets import CLIENT_SURFACE_TOOLSETS
    enabled = getattr(agent, "enabled_toolsets", None)
    if enabled is not None:
        enabled = [*enabled, *CLIENT_SURFACE_TOOLSETS]
    return _select_tool_names(enabled, getattr(agent, "disabled_toolsets", None), True)


def _drop_gated_carried_tools(merged: list, carried: set) -> list:
    """A carried tool also passes the session-level schema gates the fresh build applied
    (``browser_exec`` needs ``terminal`` in the same array), judged on the merged array."""
    from model_tools import _DYNAMIC_SCHEMA_REWRITERS
    available = {_def_name(t) for t in merged}
    return [t for t in merged if _def_name(t) not in carried or _def_name(t) not in _DYNAMIC_SCHEMA_REWRITERS
            or _DYNAMIC_SCHEMA_REWRITERS[_def_name(t)](t, available) is not None]


def _pinned_tool_defs(saved, fresh, keep_bytes=False) -> tuple:
    """(pinned, same_code, definitions) for a tools[] pin, either shape: a pin from the SAME code keeps
    its bytes; otherwise (``hermes update``, a legacy name list) each tool takes this process's
    definition (``fresh``, else its registry schema)."""
    pinned, version = (saved.get("tools") or [], saved.get("version")) if isinstance(saved, dict) else (saved, None)
    if not isinstance(pinned, list):
        return [], False, []
    from tools.registry import registry
    same_code = version is not None and version == tool_pin_version()

    def _pinned_def(item):
        name = item if isinstance(item, str) else _def_name(item) if isinstance(item, dict) else None
        if isinstance(item, dict) and (same_code or keep_bytes):
            return item
        if name in fresh:
            return fresh[name]
        entry = registry.get_entry(name) if name else None
        return {"type": "function", "function": {**entry.schema, "name": entry.name}} if entry else None

    return pinned, same_code, [d for d in map(_pinned_def, pinned) if d]


def pinned_tool_defs(saved) -> list:
    """``saved``'s tool definitions as a fresh agent restores them, for an estimate with no agent built:
    pinned definitions keep their bytes, a legacy name list takes each tool's current registry schema."""
    return _pinned_tool_defs(saved, {}, keep_bytes=True)[2]


def restore_agent_tool_prefix(agent, saved) -> bool:
    """Fold a freshly built agent's ``tools`` onto the session's pin; True if changed.
    A fresh AIAgent (gateway cache eviction, ``--resume`` in a new process, a surface hop) has no
    predecessor to preserve, so the pin stands in. Pinned by the SAME code, a tool still available
    here keeps its pinned BYTES, whatever this process derives for it (tool_search's per-surface
    catalog, per-surface dynamic parameters, the ``-q`` footprint): tools[] heads every request.
    Pinned by other code (``hermes update``, a legacy name list) a tool's contract may have moved,
    so each takes its current definition. A pinned tool this process did not build is carried
    only while its toolset config allows it here; deregistered tools drop, new tools append."""
    pinned, same_code, pinned_defs = _pinned_tool_defs(saved, {_def_name(t): t for t in _agent_tool_defs(agent)})
    if not pinned:
        return False
    from tools.registry import registry
    fresh_defs = _agent_tool_defs(agent)
    fresh = {_def_name(t): t for t in fresh_defs}
    registered_names = {entry.name for entry in registry.get_all_entries()}
    carried = {_def_name(d) for d in pinned_defs if _def_name(d) not in fresh and _def_name(d) in registered_names}
    carried &= _config_permitted_names(agent) if carried else set()
    merged = [d for d in pinned_defs if _def_name(d) in fresh or _def_name(d) in carried]
    pinned_names = {_def_name(d) for d in merged}
    merged.extend(t for t in fresh_defs if _def_name(t) not in pinned_names)
    merged = _drop_gated_carried_tools(merged, carried)
    merged_names = {_def_name(t) for t in merged}
    _reinject_authorized_dynamic_tools(agent, merged, merged_names)
    merged, merged_names = _drop_session_tools(agent, merged, merged_names)
    changed = merged != fresh_defs
    if changed:
        with _agent_tools_lock:
            agent.tools = merged
            agent.valid_tool_names = merged_names
    if not same_code or merged != list(pinned):
        persist_agent_tool_names(agent)
    return changed


def _merge_preserving_prefix(current_defs: list, new_defs: list, registered_names: set) -> tuple[list, set]:
    """Fold a fresh tool snapshot into a live one without moving existing bytes. Ordered by
    ``current_defs`` (the cached request prefix): a name already sent keeps its slot and the
    schema bytes already sent; a name only in the live list is kept if still registered
    (``check_fn`` flapped), else dropped; a name only in the fresh list is appended at the tail.

    A refreshed schema is not substituted in place. Tool JSON is rendered ahead of the
    conversation, so one changed description means the next prompt no longer starts with the
    previous one and a local prefix cache reads the whole prompt again. An explicit reload
    (``preserve_prefix`` off) still takes the fresh schemas.

    The carried entry freezes the whole schema (description and parameters), not just prose: if a
    server changes a live tool's input schema under the same name, the model keeps calling with
    the old shape and the server rejects those calls until ``/reload-mcp`` or compaction.

    The bridge tools are synthesized by ``assemble_tool_defs``, never registered, and a fresh
    snapshot omits them when the deferred set shrinks under the activation threshold: they
    keep their slot regardless (search reads the live catalog at dispatch)."""
    from tools.tool_search_catalog import BRIDGE_TOOL_NAMES
    fresh = {_def_name(entry): entry for entry in new_defs if _def_name(entry)}
    merged = []
    for entry in current_defs:
        name = _def_name(entry)
        replacement = fresh.pop(name, None)
        if replacement is not None or name in BRIDGE_TOOL_NAMES or (name and name in registered_names):
            merged.append(entry)
    merged.extend(fresh.values())
    return merged, {_def_name(t) for t in merged}


def _reinject_authorized_dynamic_tools(agent, tools_list: list, name_set: set) -> None:
    """``message_agent`` is injected by an auth gate, never registered, so a registry-derived
    rebuild drops it. Scrub any stale copy from the STAGED pair and re-add it only when the live
    gate re-authorizes, so the publisher exposes a coherent ``(tools, valid_tool_names)``."""
    from tools.bot_mode_dm import MESSAGE_AGENT_TOOL_NAME, message_agent_authorized, message_agent_tool_schema

    tools_list[:] = [entry for entry in tools_list if _def_name(entry) != MESSAGE_AGENT_TOOL_NAME]
    name_set.discard(MESSAGE_AGENT_TOOL_NAME)
    if message_agent_authorized(agent):
        tools_list.append(message_agent_tool_schema())
        name_set.add(MESSAGE_AGENT_TOOL_NAME)


def _reinject_post_build_tools(agent, tools_list: list, name_set: set) -> set:
    """Append memory-provider and context-engine tools onto the caller's staged ``tools_list``
    / ``name_set`` (never the live agent attributes), mirroring ``agent_init``'s post-build
    injection. Idempotent and fail-soft. Returns the context-engine routing names THIS rebuild
    appended: a name already owned by a registry/plugin tool is not claimed, matching agent_init."""
    def _add(schema) -> bool:
        name = schema.get("name", "") if isinstance(schema, dict) else ""
        if not name or name in name_set:
            return False
        tools_list.append({"type": "function", "function": schema})
        name_set.add(name)
        return True

    def _schema_getter(attr: str, method: str):
        getter = getattr(getattr(agent, attr, None) or None, method, None)
        return getter if callable(getter) else None

    enabled = getattr(agent, "enabled_toolsets", None)
    try:
        get_mem_schemas = _schema_getter("_memory_manager", "get_all_tool_schemas")
        if get_mem_schemas is not None:
            from agent.memory_manager_tools import memory_provider_tools_enabled  # same gate inject_memory_provider_tools uses
            if memory_provider_tools_enabled(
                    enabled, getattr(agent, "disabled_toolsets", None), memory_tool_present="memory" in name_set):
                for schema in get_mem_schemas():
                    _add(schema)
    except Exception:
        logger.debug("Memory-provider tool re-injection skipped", exc_info=True)

    # The `context_engine` toolset is intentionally empty, so lcm_* tools exist only via this
    # append. Honor the enabled_toolsets gate agent_init uses, or a restricted-toolset platform
    # would re-leak tools the build excluded.
    # See #5544.
    staged_engine_names: set = set()
    try:
        get_schemas = _schema_getter("context_compressor", "get_tool_schemas")
        if (enabled is None or "context_engine" in enabled) and get_schemas is not None:
            # Claim the routing name only when WE appended the schema.
            staged_engine_names.update(s["name"] for s in get_schemas() if _add(s))
    except Exception:
        logger.debug("Context-engine tool re-injection skipped", exc_info=True)
    return staged_engine_names