#!/usr/bin/env python3
""

import json
import os
import sys
import time
from typing import Any

sys.path.insert(0, os.path.dirname(os.path.realpath(__file__)))
import acp_lib  # after the script's own dir is on the path

DEFAULT_PROMPT = "Reply with exactly the word PONG and nothing else. Do not use any tools."
STREAM = "acp:events"


def summarize(updates):
    ""
    kinds: dict[str, int] = {}
    for u in updates:
        k = acp_lib.update_kind(u) or u.get("method")
        kinds[k] = kinds.get(k, 0) + 1
    return kinds, "".join(acp_lib.chunk_text(u) for u in updates)


def parse_args(argv, cwd):
    ""
    o = {
        "agent": "claude",
        "cwd": cwd,
        "timeout": 120.0,
        "ledger": True,
        "mode": None,
        "borrow": True,
        "text": False,
    }
    rest = []
    i = 0
    while i < len(argv):
        a = argv[i]
        if a in ("--agent", "--cwd", "--timeout", "--mode") and i + 1 < len(argv):
            v = argv[i + 1]
            if a == "--agent":
                o["agent"] = v
            elif a == "--cwd":
                o["cwd"] = os.path.expanduser(v)
            elif a == "--mode":
                o["mode"] = v
            else:
                o["timeout"] = float(v)
            i += 2
            continue
        if a == "--no-ledger":
            o["ledger"] = False
        elif a == "--no-borrow":
            o["borrow"] = False
        elif a == "--text":
            o["text"] = True
        else:
            rest.append(a)
        i += 1
    o["prompt"] = " ".join(rest) or DEFAULT_PROMPT
    return o


def main(argv):
    if "--selftest" in argv:
        return selftest()
    o = parse_args(argv, os.getcwd())
    agent, cwd, timeout, ledger_on, prompt = (
        o["agent"],
        o["cwd"],
        o["timeout"],
        o["ledger"],
        o["prompt"],
    )
    cmd, want_mode, auth, refs = acp_lib.agent_spec(agent)
    if o["mode"]:
        want_mode = o["mode"]
    if not o["borrow"]:
        refs = acp_lib.literal_refs(refs)
    ledger = None
    if ledger_on:
        import redis

        ledger = redis.Redis(decode_responses=True)
    updates, requests = [], []

    def led(**f):
        if ledger:
            try:
                ledger.xadd(
                    STREAM,
                    {"seat": "probe", **{k: str(v)[:300] for k, v in f.items()}},
                    maxlen=2000,
                    approximate=True,
                )
            except Exception as e:  # noqa: BLE001 -- a ledger outage must not end the probe
                print(f"acp-probe: ledger write failed: {e!r}", file=sys.stderr)

    def on_message(kind, msg):
        if kind == "notification":
            updates.append(msg)
        elif kind in ("permission", "request"):
            requests.append(msg.get("method"))
        led(kind=kind, update=acp_lib.update_kind(msg) if isinstance(msg, dict) else "")

    c = acp_lib.Client(cmd, cwd, on_message=on_message, env=acp_lib.resolve_env(refs))
    out: dict[str, Any] = {"agent": agent, "cmd": cmd, "cwd": cwd}
    t = time.time()
    init = c.initialize(name="thomas-acp-probe")
    out["initialize_s"] = round(time.time() - t, 2)
    res = init.get("result") or {}
    out["initialize"] = init.get("error") or {
        "protocolVersion": res.get("protocolVersion"),
        "loadSession": (res.get("agentCapabilities") or {}).get("loadSession"),
    }
    sid = None
    if auth and not init.get("error"):
        a = c.authenticate(auth)
        out["authenticate"] = a.get("error") or auth
        if a.get("error"):
            init = a
    if not init.get("error"):
        t = time.time()
        new = c.call("session/new", acp_lib.session_params(cwd, []), 60)
        out["session_new_s"] = round(time.time() - t, 2)
        sid = (new.get("result") or {}).get("sessionId")
        out["session_new"] = new.get("error") or {"sessionId": sid}
        if sid:
            offered = (new.get("result") or {}).get("modes")
            mode, err = c.enforce_mode(sid, want_mode, offered)
            out["mode"] = err or mode or "none"
            if err:
                sid = None
    if sid:
        t = time.time()
        resp = c.call(
            "session/prompt",
            {"sessionId": sid, "prompt": [{"type": "text", "text": prompt}]},
            timeout,
        )
        out["prompt_s"] = round(time.time() - t, 2)
        out["prompt"] = resp.get("error") or resp.get("result")
    out["updates"], text = summarize(updates)
    out["agent_text"] = text[:1000]
    out["agent_requests"] = requests
    ok = bool(sid) and isinstance(out.get("prompt"), dict) and "stopReason" in out["prompt"]
    if not ok:
        out["stderr_tail"] = c.stderr_tail[-15:]
    led(
        kind="summary",
        ok=ok,
        agent=agent,
        stop=(out.get("prompt") or {}).get("stopReason", ""),
        prompt_s=out.get("prompt_s", ""),
        text=text[:200],
    )
    c.close()
    print(render(out, text, ok, o["text"]))
    return 0 if ok else 1


def render(out, text, ok, text_mode):
    ""
    if text_mode and ok:
        return text.strip() or "(the assistant replied with no text)"
    return json.dumps(out, indent=1)


def selftest():
    chunk = {
        "jsonrpc": "2.0",
        "method": "session/update",
        "params": {
            "sessionId": "s",
            "update": {
                "sessionUpdate": "agent_message_chunk",
                "content": {"type": "text", "text": "PO"},
            },
        },
    }
    tool = {
        "jsonrpc": "2.0",
        "method": "session/update",
        "params": {"sessionId": "s", "update": {"sessionUpdate": "tool_call", "title": "Read"}},
    }
    opts = {
        "options": [
            {"optionId": "a1", "kind": "allow_once"},
            {"optionId": "r1", "kind": "reject_once"},
        ]
    }
    kinds, text = summarize([chunk, tool, chunk])
    cases = [
        ("route-response", acp_lib.route({"jsonrpc": "2.0", "id": 1, "result": {}}) == "response"),
        (
            "route-error-response",
            acp_lib.route({"jsonrpc": "2.0", "id": 1, "error": {"code": 1}}) == "response",
        ),
        (
            "route-request",
            acp_lib.route({"jsonrpc": "2.0", "id": 7, "method": "session/request_permission"})
            == "request",
        ),
        ("route-notification", acp_lib.route(chunk) == "notification"),
        ("route-junk", acp_lib.route({"id": 3}) == "junk" and acp_lib.route("x") == "junk"),
        (
            "update-kind",
            acp_lib.update_kind(chunk) == "agent_message_chunk"
            and acp_lib.update_kind(tool) == "tool_call",
        ),
        (
            "chunk-text-only-from-message-chunks",
            acp_lib.chunk_text(chunk) == "PO" and acp_lib.chunk_text(tool) == "",
        ),
        (
            "permission-picks-reject-never-allow",
            acp_lib.reject_policy(opts) == {"outcome": {"outcome": "selected", "optionId": "r1"}},
        ),
        (
            "permission-cancels-when-no-reject",
            acp_lib.reject_policy({"options": [{"optionId": "a", "kind": "allow_always"}]})
            == {"outcome": {"outcome": "cancelled"}},
        ),
        (
            "permission-cancels-on-empty",
            acp_lib.reject_policy({}) == {"outcome": {"outcome": "cancelled"}},
        ),
        (
            "reject-prefers-once-over-always",
            acp_lib.reject_policy(
                {
                    "options": [
                        {"optionId": "ra", "kind": "reject_always"},
                        {"optionId": "ro", "kind": "reject_once"},
                    ]
                }
            )
            == {"outcome": {"outcome": "selected", "optionId": "ro"}},
        ),
        (
            "allow-never-picks-allow-always",
            acp_lib.choose_option(
                [
                    {"optionId": "aa", "kind": "allow_always"},
                    {"optionId": "ro", "kind": "reject_once"},
                ],
                "allow",
            )
            == {"outcome": {"outcome": "selected", "optionId": "ro"}},
        ),
        (
            "allow-picks-allow-once",
            acp_lib.choose_option(opts["options"], "allow")
            == {"outcome": {"outcome": "selected", "optionId": "a1"}},
        ),
        (
            "summarize-counts-and-joins",
            kinds == {"agent_message_chunk": 2, "tool_call": 1} and text == "POPO",
        ),
    ]
    po = parse_args(
        ["--agent", "gemini", "--mode", "default", "--no-borrow", "--no-ledger", "hi", "there"],
        "/w",
    )
    cases += [
        (
            "parse-args-mode-borrow-ledger-and-prompt",
            (po["agent"], po["mode"], po["borrow"], po["ledger"], po["prompt"], po["cwd"])
            == ("gemini", "default", False, False, "hi there", "/w"),
        ),
        (
            "parse-args-defaults-keep-declared-mode-and-borrowing",
            (
                parse_args([], "/w")["mode"],
                parse_args([], "/w")["borrow"],
                parse_args([], "/w")["ledger"],
            )
            == (None, True, True),
        ),
        (
            "text-flag-is-parsed-and-is-not-part-of-the-prompt",
            (parse_args(["--text", "hi"], "/w")["text"], parse_args(["--text", "hi"], "/w")["prompt"],
             parse_args(["hi"], "/w")["text"]) == (True, "hi", False),
        ),
        (
            "text-mode-success-prints-the-whole-reply-and-no-json",
            render({"model_usage": [1]}, "x" * 1500 + "END", True, True) == "x" * 1500 + "END",
        ),
        (
            "text-mode-failure-still-prints-the-json-diagnostic",
            '"stderr_tail"' in render({"stderr_tail": ["boom"]}, "", False, True),
        ),
        ("default-mode-prints-json", render({"a": 1}, "hi", True, False).startswith("{")),
    ]
    fails = [n for n, ok in cases if not ok]
    for n, ok in cases:
        print(f"  {'PASS' if ok else 'FAIL'}  {n}")
    print(
        f"acp-probe --selftest RESULT: {'PASS' if not fails else 'FAIL'} -- {len(cases) - len(fails)}/{len(cases)} cases"
        + ("" if not fails else f" (failed: {','.join(fails)})")
    )
    return 0 if not fails else 1


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))
