#!/usr/bin/env python3
"""WS test-client harness (docs/13-testing.md §13.2).

Connects to the REAL running gateway and drives a turn, printing every
frame. This is how we empirically confirm the exact frame shapes before /
while building the Kotlin client.

Usage::

    hermes gateway &                      # with the android plugin
    python gateway-plugin/tests/ws_probe.py --token <ANDROID_TOKEN> \
        --send "hello"

Options:
    --url        ws://host:port/ws   (default ws://127.0.0.1:8790/ws)
    --token      ANDROID_TOKEN       (default: $ANDROID_TOKEN)
    --device     device_id           (default: probe-<rand>)
    --send TEXT  send this message after pairing (default: "hello")
    --timeout S  seconds to wait for the final reply (default 120)
    --authfail   expect an auth rejection (wrong token) and exit 0 on it
"""

import argparse
import asyncio
import json
import os
import sys
import time
import uuid

try:
    import websockets
except ImportError:  # pragma: no cover
    sys.stderr.write("websockets is required (hermes core dep); run inside the hermes venv\n")
    raise


def _print_frame(raw):
    try:
        data = json.loads(raw)
    except (json.JSONDecodeError, TypeError):
        print(f"  <- {raw!r}")
        return None
    ftype = data.get("type", "?")
    chat = data.get("chat_id")
    fid = data.get("id")
    payload = data.get("payload", {})
    # Compact one-line summary + full payload for the interesting frames.
    extra = ""
    if ftype == "message":
        text = (payload.get("text") or "")
        extra = f" role={payload.get('role')} id={payload.get('message_id')} text={text[:120]!r}"
        if payload.get("reasoning"):
            extra += f" reasoning={payload['reasoning'][:80]!r}"
    elif ftype == "message.start":
        extra = f" id={payload.get('message_id')} role={payload.get('role')}"
    elif ftype == "message.update":
        text = (payload.get("text") or "")
        extra = f" id={payload.get('message_id')} text={text[:100]!r}"
    elif ftype == "message.stop":
        text = (payload.get("final_text") or "")
        extra = f" id={payload.get('message_id')} text={text[:120]!r}"
        if payload.get("reasoning"):
            extra += f" reasoning={payload['reasoning'][:80]!r}"
    elif ftype == "tool.start":
        extra = (f" idx={payload.get('index')} name={payload.get('name')!r} "
                 f"preview={str(payload.get('preview'))[:80]!r}")
    elif ftype == "tool.progress":
        extra = f" idx={payload.get('index')} name={payload.get('name')!r} note={payload.get('note')!r}"
    elif ftype == "tool.end":
        extra = (f" idx={payload.get('index')} name={payload.get('name')!r} "
                 f"ok={payload.get('ok')} dur={payload.get('duration')}")
    elif ftype == "commentary":
        extra = f" id={payload.get('message_id')} text={(payload.get('text') or '')[:120]!r}"
    elif ftype == "hello.ack":
        extra = f" caps={payload.get('server_caps')} cursor={payload.get('sync_cursor')}"
    elif ftype == "error":
        extra = f" code={payload.get('code')} msg={payload.get('message')!r}"
    elif ftype == "typing":
        extra = f" on={payload.get('on')}"
    elif ftype == "pong":
        extra = ""
    scope = f" chat={chat}" if chat else ""
    idpart = f" id={fid}" if fid is not None else ""
    print(f"  <- {ftype}{idpart}{scope}{extra}")
    return data


async def run(args) -> int:
    url = args.url
    token = args.token
    device_id = args.device

    print(f"== ws_probe: connecting {url} device={device_id}")
    try:
        ws = await websockets.connect(url, open_timeout=10)
    except Exception as e:
        print(f"!! connect failed: {e}")
        return 2

    hello = {
        "v": 1,
        "type": "hello",
        "payload": {
            "token": token,
            "device_id": device_id,
            "device_name": "ws-probe",
            "caps": {"min_protocol": 1},
        },
    }
    await ws.send(json.dumps(hello))
    print("  -> hello")

    # First response must be hello.ack (or an auth error).
    try:
        first = await asyncio.wait_for(ws.recv(), timeout=10)
    except asyncio.TimeoutError:
        print("!! no hello.ack within 10s")
        await ws.close()
        return 3
    data = _print_frame(first)
    if data is None or data.get("type") != "hello.ack":
        if args.authfail:
            print("== auth rejected as expected")
            await ws.close()
            return 0
        print("!! expected hello.ack")
        await ws.close()
        return 4
    if args.authfail:
        print("!! expected auth rejection but got hello.ack")
        await ws.close()
        return 5

    if not args.send:
        print("== paired OK (no --send; exiting)")
        await ws.close()
        return 0

    # Drive a turn.
    msg_id = 1
    send_frame = {
        "v": 1,
        "id": msg_id,
        "type": "message.send",
        "chat_id": "android:default",
        "payload": {"text": args.send},
    }
    await ws.send(json.dumps(send_frame))
    print(f"  -> message.send id={msg_id} text={args.send!r}")

    deadline = time.time() + args.timeout
    got_final = False
    seen_final_frame = False
    while time.time() < deadline:
        try:
            raw = await asyncio.wait_for(ws.recv(), timeout=deadline - time.time())
        except asyncio.TimeoutError:
            print(f"!! timeout after {args.timeout}s waiting for final message")
            await ws.close()
            return 6
        data = _print_frame(raw)
        if data is None:
            continue
        ftype = data.get("type")
        payload = data.get("payload") or {}
        # A standalone assistant `message` (non-streaming) is immediately final.
        if ftype == "message" and payload.get("role") == "assistant":
            got_final = True
            break
        # A `message.stop` finalizes a streaming segment; the turn is done once
        # typing stops afterwards (multi-segment turns have several stops).
        if ftype == "message.stop":
            seen_final_frame = True
        if ftype == "typing" and payload.get("on") is False and seen_final_frame:
            got_final = True
            break
    await ws.close()
    if got_final:
        print("== final assistant message received")
        return 0
    print("!! no final assistant message")
    return 7


def main() -> int:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--url", default=os.getenv("ANDROID_WS_URL", "ws://127.0.0.1:8790/ws"))
    p.add_argument("--token", default=os.getenv("ANDROID_TOKEN", ""))
    p.add_argument("--device", default=f"probe-{uuid.uuid4().hex[:8]}")
    p.add_argument("--send", default="hello")
    p.add_argument("--timeout", type=float, default=120.0)
    p.add_argument("--authfail", action="store_true",
                   help="expect an auth rejection (wrong token)")
    args = p.parse_args()
    if not args.token and not args.authfail:
        p.error("--token (or $ANDROID_TOKEN) is required")
    return asyncio.run(run(args))


if __name__ == "__main__":
    sys.exit(main())