#!/usr/bin/env python3
"""Exercise every registered Mugi tool through the live mugi_backend (LLM + real prefs).

Uses one fresh conversation per tool. Prompts include find_tools when the tool is not tier-0,
and JSON argument hints produced from each tool's Pydantic schema (minimal required fields).

Environment:
  TOOL_PROBE_LIMIT   — max tools to run (0 = no limit). Default 0.
  TOOL_PROBE_OFFSET  — skip first N tools after sort (for sharding). Default 0.
  TOOL_PROBE_POLL_S  — seconds between /chat/messages polls. Default 2.
  TOOL_PROBE_ROUNDS — max poll iterations per tool. Default 100 (~200s at poll=2).
  TOOL_PROBE_DELETE — if "1" (default), DELETE each conversation after the tool turn.
  TOOL_PROBE_INCLUDE_CLI — if "1", load brew/ffmpeg/runner tools for introspection (can be slow).

Does not replace scripts/tool_smoke.py (direct StructuredTool.invoke); use both for coverage.

For **Research Mode** and **delegate_task → delegation scratch** (multi-turn), use
``scripts/chained_tool_probe_http.py`` (sends ``research_mode: true`` and passes
``parent_thread_id`` / ``delegation_job_id`` on scratch tools).

Skip list: installs, remote git push, mesh contact, and a few irreversible / peer-dependent tools.

Deliberate coverage map: see ``testing/deliberate_matrix.json`` and ``scripts/run_deliberate_matrix.py``.
"""
from __future__ import annotations

import json
import os
import subprocess
import sys
import time
import urllib.error
import urllib.request
from typing import Any, Container

_SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
_PROJECT_ROOT = os.path.dirname(_SCRIPT_DIR)
if _PROJECT_ROOT not in sys.path:
    sys.path.insert(0, _PROJECT_ROOT)


def _pgrep_backend_pids() -> list[int]:
    try:
        out = subprocess.check_output(["pgrep", "-f", "mugi_backend"], text=True).strip()
    except subprocess.CalledProcessError:
        return []
    return [int(x) for x in out.splitlines() if x.strip().isdigit()]


def _port_from_pid(pid: int) -> int | None:
    try:
        out = subprocess.check_output(["lsof", "-nP", "-p", str(pid)], text=True, stderr=subprocess.DEVNULL)
    except subprocess.CalledProcessError:
        return None
    for line in out.splitlines():
        if "TCP" in line and "LISTEN" in line:
            part = line.split()[-2]
            if ":" in part:
                try:
                    return int(part.rsplit(":", 1)[-1])
                except ValueError:
                    continue
    return None


def _pick_http_backend_pid() -> int | None:
    for pid in _pgrep_backend_pids():
        if _port_from_pid(pid) is not None:
            return pid
    return None


def _auth_from_pid(pid: int) -> str:
    out = subprocess.check_output(["ps", "eww", "-p", str(pid)], text=True)
    for tok in out.split():
        if tok.startswith("MUGI_AUTH_TOKEN="):
            return tok.split("=", 1)[1]
    raise SystemExit(f"No MUGI_AUTH_TOKEN in environment for pid {pid}")


def http_json(method: str, url: str, headers: dict, body: dict | None = None, timeout: float = 120) -> tuple[int, object]:
    data = None
    hdrs = {k: v for k, v in headers.items()}
    if body is not None:
        data = json.dumps(body).encode("utf-8")
        hdrs.setdefault("Content-Type", "application/json")
    req = urllib.request.Request(url, data=data, headers=hdrs, method=method)
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            raw = resp.read().decode("utf-8", errors="replace")
            code = resp.getcode()
    except urllib.error.HTTPError as e:
        raw = e.read().decode("utf-8", errors="replace")
        code = e.code
    try:
        parsed = json.loads(raw) if raw.strip() else {}
    except json.JSONDecodeError:
        parsed = raw
    return code, parsed


def _workspace_abs() -> str:
    try:
        from core import get_workspace_path

        return os.path.abspath(os.path.expanduser(get_workspace_path() or "~/.mugi/workspace"))
    except Exception:
        return os.path.abspath(os.path.expanduser("~/.mugi/workspace"))


def _flatten_tool_names() -> list[str]:
    from tools import _CLI_TOOL_NAMES, _TOOL_NAMES_BY_MODULE

    names: set[str] = set()
    for _mod, lst in _TOOL_NAMES_BY_MODULE.items():
        for n in lst:
            names.add(n)
    names |= set(_CLI_TOOL_NAMES)
    return sorted(names)


# Keep in sync with agent_loop._RESEARCH_PIPELINE_ONLY_TOOLS — not loaded in tools_by_name when research_mode is off.
_RESEARCH_PIPELINE_ONLY = frozenset({
    "analyze_multiple_sources",
    "append_research_scratch",
    "save_research_findings",
    "generate_final_research_report",
    "suggest_follow_up_queries",
    "reflect_research_state",
    "plan_research_subquestions",
})

# Tools we never auto-probe via HTTP (manual / environment / irreversible).
HTTP_PROBE_SKIP: dict[str, str] = {
    "mugi_brew_install": "installs packages; confirm in UI",
    "install_cli_package": "creates venv + pip network",
    "git_project_push": "mutates remote",
    "contact_peer": "requires live mesh peer",
    "mesh_send_files": "requires live mesh peer + paths",
    "download_video_yt_dlp": "long-running + ToS-sensitive",
    "capture_network_traffic": "browser capture; heavy",
    "screenshot_webpage": "headful stack; flaky in CI",
    "burn_subtitles": "ffmpeg mutate; needs crafted inputs",
    "images_to_video": "heavy CLI",
    "merge_pdfs": "needs two PDF paths",
    "split_pdf": "mutates documents",
    "convert_video": "heavy unless tiny fixture exists",
    "trim_media": "needs media path",
    "concat_media": "needs multiple inputs",
    "merge_audio_video": "needs inputs",
    "remove_audio": "mutates media",
    "extract_subtitles": "needs media with subs",
    "split_media": "mutates media",
    "rotate_video": "mutates media",
    "change_speed": "mutates media",
    "normalize_audio": "mutates media",
    "extract_frames_interval": "many files",
    "diagnose_plugin_sandbox_migration": "niche migration helper",
}

for _rn in _RESEARCH_PIPELINE_ONLY:
    HTTP_PROBE_SKIP.setdefault(
        _rn,
        "requires Research Mode (chat UI); pipeline tools stripped when research_mode is false",
    )
for _delegation_scratch_tool in ("append_delegation_scratch", "read_delegation_scratch"):
    HTTP_PROBE_SKIP.setdefault(
        _delegation_scratch_tool,
        "single-turn probe cannot set parent thread/job — run scripts/chained_tool_probe_http.py",
    )


def _minimal_args_from_schema(tool_obj: Any, workspace: str) -> dict[str, Any]:
    """Build a small dict covering required fields; optional fields omitted."""
    schema_obj = getattr(tool_obj, "args_schema", None)
    if schema_obj is None:
        return {}
    try:
        if hasattr(schema_obj, "model_json_schema"):
            schema = schema_obj.model_json_schema()
        elif hasattr(schema_obj, "schema"):
            schema = schema_obj.schema()
        else:
            return {}
    except Exception:
        return {}
    props: dict[str, Any] = dict(schema.get("properties") or {})
    required: list[str] = list(schema.get("required") or [])
    out: dict[str, Any] = {}

    def fill(key: str) -> None:
        if key in out:
            return
        p = props.get(key) or {}
        typ = p.get("type")
        key_l = key.lower()
        if typ == "string":
            if "url" in key_l or key_l == "webhook_url":
                out[key] = "https://example.com/"
            elif "path" in key_l or "file" in key_l or key_l in ("dir_path", "filepath", "root_path"):
                if "chunk" in key_l or "ocr" in key_l or "pdf" in key_l or "image" in key_l:
                    out[key] = os.path.join(workspace, "TOOLS.md")
                else:
                    out[key] = "."
            elif key_l in ("prompt", "task_description", "text", "content", "topic", "query"):
                if "task" in key_l or key_l == "prompt":
                    out[key] = "Reply with one word: PROBE_OK"
                elif key_l == "query":
                    out[key] = "mugi"
                elif key_l == "topic":
                    out[key] = "probe_topic"
                else:
                    out[key] = "probe"
            elif key_l == "schedule":
                out[key] = "in 365 days"
            elif key_l == "tool_name":
                out[key] = "get_current_time"
            elif key_l == "tool_args":
                out[key] = "{}"
            elif key_l == "command" and key == "command":
                out[key] = "true"
            else:
                out[key] = "probe"
        elif typ == "integer":
            if "hours" in key_l:
                out[key] = 24
            elif "max" in key_l or "limit" in key_l or "results" in key_l:
                out[key] = 3
            else:
                out[key] = 1
        elif typ == "boolean":
            out[key] = False
        elif typ == "number":
            out[key] = 1.0
        elif typ == "array":
            if key_l == "tasks":
                out[key] = [
                    {"task_description": "Reply one word: A", "agent_type": "general"},
                ]
            elif "sources" in key_l or "urls" in key_l:
                out[key] = ["https://example.com/"]
            else:
                out[key] = []
        elif typ == "object":
            out[key] = {}
        else:
            out[key] = ""

    for r in required:
        fill(r)
    # Tools that break on empty required object
    if not required and getattr(tool_obj, "name", "") in (
        "list_directory",
        "read_file",
        "verify_file_exists",
    ):
        if getattr(tool_obj, "name", "") == "list_directory":
            out["dir_path"] = "."
        elif getattr(tool_obj, "name", "") == "read_file":
            out["filepath"] = "TOOLS.md"
        elif getattr(tool_obj, "name", "") == "verify_file_exists":
            out["file_path"] = "TOOLS.md"
    return out


def _resolve_placeholders(obj: Any, workspace: str) -> Any:
    if isinstance(obj, str):
        return obj.replace("{{WORKSPACE}}", workspace)
    if isinstance(obj, dict):
        return {k: _resolve_placeholders(v, workspace) for k, v in obj.items()}
    if isinstance(obj, list):
        return [_resolve_placeholders(v, workspace) for v in obj]
    return obj


def _build_prompt(tool_name: str, args_hint: dict[str, Any], tier0: Container[str]) -> str:
    args_json = json.dumps(args_hint, ensure_ascii=False)
    find_step = ""
    if tool_name not in tier0:
        find_step = (
            f"Step A: Call find_tools once with query {tool_name!r} and max_results 8. "
        )
    return (
        f"QA single-tool probe (harmless). {find_step}"
        f"Step B: Call the tool {tool_name!r} exactly once. "
        f"Use arguments as close as possible to this JSON (fix types only if the schema requires it): {args_json}. "
        "Do not call any other tools. "
        "After the tool returns, reply exactly one line starting with PROBE_RESULT= "
        "followed by ok or fail and at most 8 words."
    )


def _tool_rows(msgs: list) -> list[dict]:
    rows: list[dict] = []
    for m in msgs:
        if (m.get("role") or "").strip() != "tool":
            continue
        raw = m.get("content") or ""
        if not isinstance(raw, str):
            continue
        try:
            meta = json.loads(raw)
        except json.JSONDecodeError:
            continue
        if isinstance(meta, dict) and meta.get("tool_name"):
            rows.append(meta)
    return rows


def _ran_target_tool(rows: list[dict], target: str) -> bool:
    for r in rows:
        if (r.get("tool_name") or "").strip() == target:
            return True
    return False


def main() -> int:
    from tools import TIER_0_TOOLS, get_tools
    from tools.tool_index import set_tool_catalog

    pid = _pick_http_backend_pid()
    if not pid:
        print("No listening mugi_backend (pgrep -f mugi_backend)", file=sys.stderr)
        return 2
    port = _port_from_pid(pid)
    if not port:
        print(f"No LISTEN port for pid {pid}", file=sys.stderr)
        return 2
    token = _auth_from_pid(pid)
    base = f"http://127.0.0.1:{port}"
    h = {"X-Auth-Token": token}

    # Default: match typical parent prefs (CLI tools off unless include_cli_tools pref).
    # Opt in: TOOL_PROBE_INCLUDE_CLI=1 — can be slow on first import.
    include_cli = os.environ.get("TOOL_PROBE_INCLUDE_CLI", "").strip().lower() in ("1", "true", "yes")
    discovery = list(get_tools(include_media=True, include_cli=include_cli))
    if include_cli:
        try:
            from tools.cli import get_cli_tools

            discovery = discovery + list(get_cli_tools())
        except Exception:
            pass
    set_tool_catalog(discovery)
    by_name = {getattr(t, "name", ""): t for t in discovery if getattr(t, "name", None)}
    workspace = _workspace_abs()

    names = _flatten_tool_names()
    limit = int(os.environ.get("TOOL_PROBE_LIMIT", "0") or "0")
    offset = int(os.environ.get("TOOL_PROBE_OFFSET", "0") or "0")
    if offset:
        names = names[offset:]
    if limit > 0:
        names = names[:limit]

    tier0 = set(TIER_0_TOOLS)
    poll_s = float(os.environ.get("TOOL_PROBE_POLL_S", "2") or "2")
    max_rounds = int(os.environ.get("TOOL_PROBE_ROUNDS", "100") or "100")
    do_delete = os.environ.get("TOOL_PROBE_DELETE", "1") == "1"

    print(
        f"backend pid={pid} port={port} tools_planned={len(names)} workspace={workspace} "
        f"include_cli={include_cli}\n",
        flush=True,
    )

    summary: list[tuple[str, str, str]] = []

    for idx, tname in enumerate(names, start=1):
        skip_reason = HTTP_PROBE_SKIP.get(tname)
        if skip_reason:
            print(f"[{idx}/{len(names)}] SKIP {tname} — {skip_reason}", flush=True)
            summary.append((tname, "SKIP", skip_reason))
            continue

        tool_obj = by_name.get(tname)
        if tool_obj is None:
            print(
                f"[{idx}/{len(names)}] SKIP {tname} — not in discovery set (prefs/gates; try TOOL_PROBE_INCLUDE_CLI=1)",
                flush=True,
            )
            summary.append((tname, "SKIP", "not loaded"))
            continue

        raw_args = _minimal_args_from_schema(tool_obj, workspace)
        args_hint = _resolve_placeholders(raw_args, workspace)
        prompt = _build_prompt(tname, args_hint, tier0)

        code, conv = http_json("POST", f"{base}/conversations/new", h, {}, timeout=60)
        if code != 200 or not isinstance(conv, dict) or "thread_id" not in conv:
            print(f"[{idx}/{len(names)}] FAIL {tname} — new conversation {code} {conv!r}", flush=True)
            summary.append((tname, "FAIL", "new conv"))
            continue
        tid = conv["thread_id"]

        t0 = time.monotonic()
        code, ack = http_json("POST", f"{base}/send", h, {"text": prompt, "thread_id": tid}, timeout=30)
        if code != 200:
            print(f"[{idx}/{len(names)}] FAIL {tname} — /send {code} {ack!r}", flush=True)
            summary.append((tname, "FAIL", "send"))
            if do_delete:
                http_json("DELETE", f"{base}/conversations/{tid}", h, timeout=30)
            continue

        msgs: list = []
        assistant_ok = False
        for _ in range(max_rounds):
            time.sleep(poll_s)
            c2, msgs = http_json("GET", f"{base}/chat/messages?thread_id={tid}", h, timeout=60)
            if c2 != 200 or not isinstance(msgs, list):
                continue
            if msgs and (msgs[-1].get("role") or "").strip() == "assistant":
                assistant_ok = True
                break
        elapsed = time.monotonic() - t0
        rows = _tool_rows(msgs) if isinstance(msgs, list) else []
        ran = _ran_target_tool(rows, tname)

        status = "OK" if (assistant_ok and ran) else "WEAK"
        if not assistant_ok:
            status = "TIMEOUT"
        elif not ran:
            status = "NO_TOOL_ROW"

        preview = ""
        if rows:
            last = rows[-1]
            preview = str(last.get("result_preview") or "")[:160].replace("\n", " ")
        print(
            f"[{idx}/{len(names)}] {status:10} {tname:40} {elapsed:6.1f}s  "
            f"tool_rows={len(rows)} last_preview={preview!r}",
            flush=True,
        )
        summary.append((tname, status, preview[:80]))

        if do_delete:
            http_json("DELETE", f"{base}/conversations/{tid}", h, timeout=30)

    from collections import Counter

    c = Counter(s for _, s, _ in summary)
    print("\n=== summary counts ===", flush=True)
    for k, v in sorted(c.items(), key=lambda x: (-x[1], x[0])):
        print(f"  {k}: {v}", flush=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
