# agent_loop_tools.py - OpenAI tool schema conversion and plugin manager wiring extracted from agent_loop.
from __future__ import annotations

import logging
import os
import time

from tools import get_tools, get_tools_tiered

_log = logging.getLogger(__name__)


_plugin_manager = None
_knowledge_presence_cache = False
_knowledge_presence_last_scan = 0.0
_KNOWLEDGE_PRESENCE_RESCAN_SECONDS = 60.0


def _tool_to_openai_schema(tool):
    """Convert a LangChain-style tool to OpenAI function schema, or None if conversion fails."""
    name = getattr(tool, "name", None) or "unknown"
    description = getattr(tool, "description", None) or ""
    params = {"type": "object", "properties": {}, "additionalProperties": False}
    schema = getattr(tool, "args_schema", None)
    if schema is not None:
        try:
            if isinstance(schema, dict):
                # Subprocess/sandbox plugins send JSON Schema from Pydantic model_json_schema().
                raw = dict(schema)
            elif hasattr(schema, "model_json_schema"):
                raw = schema.model_json_schema()
            elif hasattr(schema, "schema"):
                raw = schema.schema()
            else:
                raw = {}
            if isinstance(raw, dict):
                params = raw.copy()
                params.pop("title", None)
                params.setdefault("type", "object")
                params.setdefault("properties", {})
        except Exception as exc:
            _log.warning("OpenAI schema conversion failed for tool %s: %s", name, exc)
            return None
    return {
        "type": "function",
        "function": {
            "name": name,
            "description": description,
            "parameters": params,
        },
    }


def set_plugin_manager(pm):
    """Set the PluginManager instance so plugin tools and hooks are available to the agent loop."""
    global _plugin_manager
    _plugin_manager = pm


def _get_plugin_manager():
    """Return the current PluginManager instance (or None)."""
    return _plugin_manager


def _knowledge_corpora_available() -> bool:
    global _knowledge_presence_cache, _knowledge_presence_last_scan
    now = time.time()
    if (now - _knowledge_presence_last_scan) < _KNOWLEDGE_PRESENCE_RESCAN_SECONDS:
        return _knowledge_presence_cache
    _knowledge_presence_last_scan = now
    try:
        from core import get_workspace_path

        workspace = (get_workspace_path() or "").strip() or os.path.expanduser("~/.mugi/workspace")
    except Exception:
        workspace = os.path.expanduser("~/.mugi/workspace")
    root = os.path.join(os.path.expanduser(workspace), "knowledge")
    present = False
    try:
        for entry in os.listdir(root):
            if entry.startswith("."):
                continue
            corpus_dir = os.path.join(root, entry)
            if not os.path.isdir(corpus_dir):
                continue
            if os.path.isfile(os.path.join(corpus_dir, "manifest.json")):
                present = True
                break
    except Exception:
        present = False
    _knowledge_presence_cache = present
    return present


# Extra tool names that should be visible in the tiered parent when the run
# originated from the Quick Panel. Keep this list small and OS-facing — the
# intent is to reduce `find_tools` round-trips for common launcher queries
# (files on disk, recents, inbox), not to expose everything.
QUICK_PANEL_TIER_MERGE_TOOL_NAMES: tuple[str, ...] = (
    "list_inbox",
    "spotlight_search",
    "get_recent_files",
    "skills_list",
)

# Dashboard composer merge: machine-diagnosis first ("why is my computer slow?",
# "what's downloading right now?", "why isn't my internet working?") plus the
# same disk/recents lookups as the Quick Panel. list_inbox/skills_list are
# intentionally dropped here to keep the schema budget flat.
DASHBOARD_TIER_MERGE_TOOL_NAMES: tuple[str, ...] = (
    "inspect_system_activity",
    "run_host_diagnostics",
    "spotlight_search",
    "get_recent_files",
)

# Dashboard schema diet: tier-0 tools dropped from the dashboard composer's
# visible schema list. The dashboard surface is machine-diagnosis + tile Q&A;
# at small local context windows (32k) the full tier-0 set (~54 schemas /
# ~15k tokens) saturates the request and the model returns empty replies.
# Excluded tools remain fully executable — they move into the find_tools
# discovery index, so the model can still reach them in two hops.
DASHBOARD_TIER_EXCLUDE_TOOL_NAMES: frozenset[str] = frozenset({
    # mesh / delegation
    "delegate_task",
    "read_delegation_scratch",
    "contact_peer",
    "list_mesh_peers",
    "mesh_send_files",
    # kanban authoring (list/show stay for the dashboard kanban tile)
    "kanban_create_task",
    "kanban_create_batch",
    "kanban_cancel_task",
    "kanban_list_runs",
    "convene_council",
    # scheduling
    "schedule_task",
    "cancel_scheduled_task",
    "list_scheduled_tasks",
    # skills
    "skills_list",
    "skill_view",
    "skill_manage",
    # task plan
    "update_task_plan",
    "clear_task_plan",
    "get_task_plan",
    # user model proposals (patch_user_model stays for user-stated facts)
    "patch_user_profile",
    "propose_user_model_update",
    "list_pending_proposals",
    # memory maintenance
    "migrate_memory_to_shards",
    # browser session capture
    "get_captured_browser_page",
    "request_browser_page",
    # clipboard
    "get_clipboard_text",
    "set_clipboard_text",
    # file writes beyond save_to_file
    "write_text_file",
    "append_to_text_file",
    # knowledge corpora
    "query_knowledge",
    "ask_studied_corpus",
    "read_knowledge_source",
    "search_memory_and_knowledge_corpora",
    "explore_knowledge_connections",
})


def get_openai_tools(
    include_media: bool = True,
    include_cli: bool = False,
    tiered: bool = False,
    quick_panel: bool = False,
    dashboard_surface: bool = False,
    core_diet: bool = False,
    working_set_names: "list[str] | tuple[str, ...] | None" = None,
):
    """Return tool schemas and lookups for execution.

    Returns:
      (openai_tool_schemas, visible_tools_by_name, full_tools_by_name, tool_index_text)
    """
    # Visible tools respect include_media (keeps default schema payload lean).
    visible_full_raw = get_tools(include_media=include_media, include_cli=include_cli)

    # Discovery/execution catalog always includes media tools so explicit find_tools
    # queries can surface OCR/PDF capabilities even when no file is attached.
    discovery_full_raw = get_tools(include_media=True, include_cli=include_cli)

    if _plugin_manager is not None:
        _plugin_manager.wait_initial_discover(timeout_s=120.0)
        plugin_tools = _plugin_manager.get_all_tools()
        visible_full_raw = list(visible_full_raw) + plugin_tools
        discovery_full_raw = list(discovery_full_raw) + plugin_tools

    mcp_tools = []
    mcp_schemas = []
    try:
        from mcp_bridge import get_mcp_discovery
        mcp_tools, mcp_schemas = get_mcp_discovery()
        mcp_tools = mcp_tools or []
        mcp_schemas = mcp_schemas or []
    except Exception:
        _log.warning("MCP discovery failed; continuing without MCP tools", exc_info=True)
        mcp_tools = []
        mcp_schemas = []
    if mcp_tools:
        # Discovery/execution includes MCP wrappers by default; visible list remains tier-controlled.
        discovery_full_raw = list(discovery_full_raw) + list(mcp_tools)

    raw = (
        get_tools_tiered(include_media=include_media, include_cli=include_cli, tier0_only=True)
        if tiered
        else list(visible_full_raw)
    )
    _diet_keep: set[str] | None = None
    if core_diet and not dashboard_surface:
        # Core diet (tool-schema subsetting): carry only the slim always-on core
        # plus this thread's earned working set. Everything dropped here stays
        # executable (full_by_name) and discoverable (find_tools / tool index),
        # and rejoins the visible set on first use. Conditional stamps below
        # (knowledge corpora, promoted plugins, web-act, operator merges) still
        # apply after the diet — they are deliberate surface decisions, not bloat.
        # Applies to both the tiered chat path and the untiered subagent path:
        # default-kit subagents otherwise carry the full catalog (182 schemas /
        # ~41k tokens — more than a 32k context cap before a single message).
        from tools import CORE_TOOLS

        _keep = set(CORE_TOOLS)
        _ws_clean = [
            s for s in (str(n or "").strip() for n in (working_set_names or [])) if s
        ]
        _keep.update(_ws_clean)
        _diet_keep = set(_keep)
        raw = [t for t in raw if (getattr(t, "name", None) or "") in _keep]
        # Working-set names outside tier-0 (media, plugins, MCP wrappers, CLI)
        # come from the discovery catalog so a resumed thread keeps its kit.
        discovery_by_name = {getattr(t, "name", None) or "": t for t in discovery_full_raw}
        seen = {(getattr(t, "name", None) or "") for t in raw}
        for _n in _ws_clean:
            _tool = discovery_by_name.get(_n)
            if _tool is not None and _n not in seen:
                raw.append(_tool)
                seen.add(_n)
    if tiered and _knowledge_corpora_available():
        visible_by_name = {getattr(t, "name", None) or "": t for t in visible_full_raw}
        seen = {(getattr(t, "name", None) or "") for t in raw}
        for _name in ("list_knowledge", "query_knowledge", "ask_studied_corpus"):
            _tool = visible_by_name.get(_name)
            if _tool is not None and _name not in seen:
                raw.append(_tool)
                seen.add(_name)
    if tiered and _plugin_manager is not None:
        # Tier-0 excludes plugins by default, but bundled/manifest-promoted tools can opt in.
        promoted = []
        for t in visible_full_raw:
            meta = getattr(t, "metadata", None) or {}
            if isinstance(meta, dict) and bool(meta.get("visible_in_tier0", False)):
                promoted.append(t)
        if promoted:
            seen = {(getattr(t, "name", None) or "") for t in raw}
            for t in promoted:
                name = getattr(t, "name", None) or ""
                if name and name not in seen:
                    raw.append(t)
                    seen.add(name)

    if tiered and dashboard_surface:
        # Dashboard schema diet: drop tier-0 tools irrelevant to the dashboard
        # surface so the schema budget fits small local context windows. The
        # excluded names stay in the discovery catalog and tool index below,
        # so find_tools can still surface and execute them.
        raw = [
            t for t in raw
            if (getattr(t, "name", None) or "") not in DASHBOARD_TIER_EXCLUDE_TOOL_NAMES
        ]

    if tiered and (quick_panel or dashboard_surface):
        # Quick Panel and Dashboard composer get an OS-facing tier-0 merge.
        # Local models do better with these tools in context than a two-hop find_tools pass.
        # The dashboard list leads with machine diagnosis (inspect_system_activity).
        merge_names = DASHBOARD_TIER_MERGE_TOOL_NAMES if dashboard_surface else QUICK_PANEL_TIER_MERGE_TOOL_NAMES
        discovery_by_name = {getattr(t, "name", None) or "": t for t in discovery_full_raw}
        seen = {(getattr(t, "name", None) or "") for t in raw}
        for _name in merge_names:
            _tool = discovery_by_name.get(_name)
            if _tool is not None and _name not in seen:
                raw.append(_tool)
                seen.add(_name)

    if tiered:
        try:
            from tools.web_act import web_act_enabled
            from tools import WEB_ACT_TOOL_NAMES

            if web_act_enabled():
                discovery_by_name = {getattr(t, "name", None) or "": t for t in discovery_full_raw}
                seen = {(getattr(t, "name", None) or "") for t in raw}
                for _name in WEB_ACT_TOOL_NAMES:
                    _tool = discovery_by_name.get(_name)
                    if _tool is not None and _name not in seen:
                        raw.append(_tool)
                        seen.add(_name)
        except Exception:
            pass

    if tiered:
        try:
            from core.config import get_config

            _cfg_t0 = (get_config() or {}).get("agent_loop_tier0") or {}
            _extra = _cfg_t0.get("merge_tool_names") or []
            if isinstance(_extra, list) and _extra:
                discovery_by_name = {getattr(t, "name", None) or "": t for t in discovery_full_raw}
                seen = {(getattr(t, "name", None) or "") for t in raw}
                for _name in _extra:
                    _n = str(_name or "").strip()
                    if not _n:
                        continue
                    _tool = discovery_by_name.get(_n)
                    if _tool is not None and _n not in seen:
                        raw.append(_tool)
                        seen.add(_n)
        except Exception:
            pass

    schemas = []
    by_name = {}
    for t in raw:
        schema = _tool_to_openai_schema(t)
        if schema is None:
            continue
        schemas.append(schema)
        by_name[getattr(t, "name", None) or "unknown"] = t
    for s in mcp_schemas:
        # When mcp_schemas_in_tiered_parent is False and tiered mode is active,
        # skip appending MCP schemas to the parent's visible tool list.
        if tiered:
            try:
                from tool_policy import get_tool_exposure_policy

                exposure_policy = get_tool_exposure_policy()
                if not exposure_policy.mcp_schemas_in_tiered_parent:
                    break
            except Exception:
                _log.warning(
                    "tool exposure policy load failed; skipping MCP schemas in tiered parent (fail-closed)",
                    exc_info=True,
                )
                break
        try:
            name = ((s or {}).get("function") or {}).get("name")
            if name and _diet_keep is not None and name not in _diet_keep:
                # Core diet on the untiered path: MCP wrappers stay executable
                # (full_by_name) and discoverable (find_tools → promotion);
                # only working-set members ride along as native schemas.
                continue
            if name and name not in by_name:
                schemas.append(s)
        except Exception:
            continue

    tool_index_text = ""
    try:
        from tools.tool_index import set_tool_catalog, build_tool_index

        # Keep find_tools catalog fresh for both parent and subagent runs.
        set_tool_catalog(list(discovery_full_raw))
        if tiered:
            visible_schema_names: set[str] = set()
            for s in schemas:
                fn = (s or {}).get("function") if isinstance(s, dict) else None
                nm = (fn or {}).get("name") if isinstance(fn, dict) else None
                if nm:
                    visible_schema_names.add(str(nm))
            tool_index_text = build_tool_index(
                list(discovery_full_raw),
                exclude_names=visible_schema_names,
            )
    except Exception:
        _log.warning(
            "tool index build failed; tiered parent will have empty discoverable index",
            exc_info=True,
        )
        tool_index_text = ""

    full_by_name = {}
    for t in discovery_full_raw:
        full_by_name[getattr(t, "name", None) or "unknown"] = t
    return schemas, by_name, full_by_name, tool_index_text


def promote_working_set_tools(
    openai_tools: list,
    tools_by_name: dict,
    visible_tool_names: set,
    full_tools_by_name: dict,
    names,
) -> list[str]:
    """Promote discovered/used tools into the live visible schema set, in place.

    Called by the agent loop at round end when the core diet is active: tools
    surfaced by find_tools or invoked by name get native schemas for every
    subsequent round of the run. Returns the names actually promoted (callers
    persist those to the per-thread working set ledger). Never raises.
    """
    promoted: list[str] = []
    try:
        for _n in sorted({str(n or "").strip() for n in (names or [])} - {""}):
            if _n in visible_tool_names:
                continue
            _t = full_tools_by_name.get(_n)
            if _t is None:
                continue
            _schema = _tool_to_openai_schema(_t)
            if _schema is None:
                continue
            openai_tools.append(_schema)
            tools_by_name[_n] = _t
            visible_tool_names.add(_n)
            promoted.append(_n)
    except Exception:
        _log.debug("working set promotion failed", exc_info=True)
    return promoted
