#!/usr/bin/env python3
"""Live backend probes that need Research Mode or multi-turn delegation context.

1) Research pipeline — POST /send with JSON body field ``research_mode: true`` so
   ``append_research_scratch`` / ``reflect_research_state`` (etc.) stay in ``tools_by_name``.

2) Delegation scratch — turn 1: ``delegate_task`` with a fixed ``delegation_job_id``;
   turn 2: ``append_delegation_scratch`` + ``read_delegation_scratch`` with explicit
   ``parent_thread_id`` (the conversation thread) and the same job id (parent agent has
   no subagent contextvars).

Env: TOOL_PROBE_POLL_S (default 2), TOOL_PROBE_ROUNDS (default 120), TOOL_PROBE_DELETE (default 1).

Run from repo root:
  /Users/you/.mugi/venv/bin/python3 -u scripts/chained_tool_probe_http.py

Matrix: ``testing/deliberate_matrix.json`` scenario ``http.chained.research_and_delegation``.
"""
from __future__ import annotations

import json
import os
import subprocess
import sys
import time
import urllib.error
import urllib.request

_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 _tool_rows(msgs: list, *, start_index: int = 0) -> list[dict]:
    rows: list[dict] = []
    for i, m in enumerate(msgs):
        if i < start_index:
            continue
        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 _tool_names_from_rows(rows: list[dict]) -> list[str]:
    return [str(r.get("tool_name") or "") for r in rows if r.get("tool_name")]


def wait_done(
    base: str,
    h: dict,
    tid: str,
    start_len: int,
    *,
    poll_s: float,
    max_rounds: int,
) -> tuple[list, bool]:
    msgs: list = []
    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 len(msgs) > start_len and msgs and (msgs[-1].get("role") or "").strip() == "assistant":
            return msgs, True
    return msgs if isinstance(msgs, list) else [], False


def main() -> int:
    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}

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

    job_id = (os.environ.get("TOOL_PROBE_DELEGATION_JOB_ID") or "chain_probe_job").strip() or "chain_probe_job"

    print(f"backend pid={pid} port={port} delegation_job_id={job_id!r}\n", flush=True)

    # --- Phase A: research_mode ---
    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("new conversation failed", code, conv, file=sys.stderr)
        return 1
    tid_r = conv["thread_id"]
    _, j0 = http_json("GET", f"{base}/chat/messages?thread_id={tid_r}", h, timeout=30)
    n0 = len(j0) if isinstance(j0, list) else 0

    research_prompt = (
        "Research-mode smoke test (harmless). "
        "Step 1: Call append_research_scratch once with topic 'http_chain_probe' "
        "and content '- probe scratch line'. "
        "Step 2: Call reflect_research_state once with original_query 'probe' and current_findings ''. "
        "Step 3: Reply one line starting with PROBE_RESEARCH= ok or fail."
    )
    code, _ = http_json(
        "POST",
        f"{base}/send",
        h,
        {"text": research_prompt, "thread_id": tid_r, "research_mode": True},
        timeout=30,
    )
    print("POST /send research_mode=True", code, flush=True)
    if code != 200:
        if do_delete:
            http_json("DELETE", f"{base}/conversations/{tid_r}", h, timeout=30)
        return 1

    msgs_r, ok_r = wait_done(base, h, tid_r, n0, poll_s=poll_s, max_rounds=max_rounds)
    rows_r = _tool_rows(msgs_r)
    names_r = _tool_names_from_rows(rows_r)
    has_append = "append_research_scratch" in names_r
    has_reflect = "reflect_research_state" in names_r
    print(
        f"Research phase: assistant_ok={ok_r} tool_calls={names_r} "
        f"(append={has_append} reflect={has_reflect})",
        flush=True,
    )
    if do_delete:
        http_json("DELETE", f"{base}/conversations/{tid_r}", h, timeout=30)

    # --- Phase B: delegate + delegation scratch ---
    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("new conversation (delegation) failed", code, conv, file=sys.stderr)
        return 1
    tid_d = conv["thread_id"]
    _, j0 = http_json("GET", f"{base}/chat/messages?thread_id={tid_d}", h, timeout=30)
    n0 = len(j0) if isinstance(j0, list) else 0

    d1 = (
        f"Call delegate_task exactly once: agent_type general, "
        f"task_description 'Reply with one word only: DELEGATE_OK', "
        f"delegation_job_id {job_id!r}. No other tools."
    )
    code, _ = http_json(
        "POST",
        f"{base}/send",
        h,
        {"text": d1, "thread_id": tid_d, "research_mode": False},
        timeout=30,
    )
    print("POST /send delegation turn1", code, flush=True)
    if code != 200:
        if do_delete:
            http_json("DELETE", f"{base}/conversations/{tid_d}", h, timeout=30)
        return 1

    msgs_d1, ok1 = wait_done(base, h, tid_d, n0, poll_s=poll_s, max_rounds=max_rounds)
    rows1 = _tool_rows(msgs_d1)
    print(f"Delegation turn1: assistant_ok={ok1} tools={_tool_names_from_rows(rows1)}", flush=True)

    _, j1 = http_json("GET", f"{base}/chat/messages?thread_id={tid_d}", h, timeout=30)
    n1 = len(j1) if isinstance(j1, list) else 0

    d2 = (
        f"Call append_delegation_scratch once with content 'Parent scratch after delegate.', "
        f"parent_thread_id {tid_d!r}, delegation_job_id {job_id!r}. "
        f"Then call read_delegation_scratch with parent_thread_id {tid_d!r}, "
        f"delegation_job_id {job_id!r}, max_chars 4000. "
        "Reply one line PROBE_DELEGATION= ok or fail."
    )
    code, _ = http_json(
        "POST",
        f"{base}/send",
        h,
        {"text": d2, "thread_id": tid_d, "research_mode": False},
        timeout=30,
    )
    print("POST /send delegation turn2", code, flush=True)
    if code != 200:
        if do_delete:
            http_json("DELETE", f"{base}/conversations/{tid_d}", h, timeout=30)
        return 1

    msgs_d2, ok2 = wait_done(base, h, tid_d, n1, poll_s=poll_s, max_rounds=max_rounds)
    # Only tools after turn-2 user message (otherwise delegate_task from turn 1 is included).
    rows2 = _tool_rows(msgs_d2, start_index=n1)
    names2 = _tool_names_from_rows(rows2)
    has_append_d = "append_delegation_scratch" in names2
    has_read_d = "read_delegation_scratch" in names2
    print(
        f"Delegation turn2: assistant_ok={ok2} tools={names2} "
        f"(append={has_append_d} read={has_read_d})",
        flush=True,
    )

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

    # Exit non-zero if phases clearly failed
    research_ok = ok_r and has_append and has_reflect
    deleg_ok = ok1 and ok2 and has_append_d and has_read_d
    print(f"\nSummary: research_phase={'OK' if research_ok else 'FAIL'} delegation_phase={'OK' if deleg_ok else 'FAIL'}", flush=True)
    return 0 if (research_ok and deleg_ok) else 1


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