#!/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")
    --upload F   M4: upload F (chunked media.upload) and attach it to the
                 message.send via media_refs
--pull-offer M4: when a media.offer arrives during the turn, pull the
                  media (chunked) and verify the byte count
    --sync C     M5: after pairing, send sync {cursor: C} and print the
                  replay + sync.done (no turn is driven)
    --fcm-token  M5: attach this FCM token to the hello payload
    --fcm-reg    M5: after pairing, send fcm.register with --fcm-token
    --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 hashlib
import json
import mimetypes
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 = ""
    elif ftype == "media.offer":
        extra = (f" media_id={payload.get('media_id')} kind={payload.get('kind')} "
                 f"mime={payload.get('mime')} size={payload.get('size')} "
                 f"file={payload.get('filename')!r} msg={payload.get('message_id')}")
    elif ftype == "media.upload.ack":
        extra = f" ok={payload.get('ok')} ref={payload.get('media_ref')}"
    elif ftype == "media.pull.end":
        extra = f" ok={payload.get('ok')}"
    elif ftype == "notification":
        extra = (f" kind={payload.get('kind')} title={payload.get('title')!r} "
                 f"body={(payload.get('body') or '')[:100]!r}")
    elif ftype == "sync":
        extra = f" cursor={payload.get('cursor')}"
    elif ftype == "sync.done":
        extra = f" cursor={payload.get('cursor')}"
    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


def _kind_for_path(path: str) -> str:
    mime, _ = mimetypes.guess_type(path)
    mime = mime or "application/octet-stream"
    if mime.startswith("image/"):
        return "image"
    if mime.startswith("video/"):
        return "video"
    if mime.startswith("audio/"):
        return "audio"
    return "document"


async def upload_file(ws, path: str, media_ref: str, next_id: int) -> int:
    """Drive media.upload.start -> binary chunks -> media.upload.end.

    Returns the next free request id; raises on a non-ack terminal frame.
    """
    data = open(path, "rb").read()
    mime, _ = mimetypes.guess_type(path)
    await ws.send(json.dumps({
        "v": 1, "id": next_id, "type": "media.upload.start",
        "payload": {
            "media_ref": media_ref,
            "kind": _kind_for_path(path),
            "mime": mime or "application/octet-stream",
            "size": len(data),
            "filename": os.path.basename(path),
        },
    }))
    print(f"  -> media.upload.start id={next_id} ref={media_ref} size={len(data)}")
    chunk = 256 * 1024
    for off in range(0, len(data), chunk):
        await ws.send(data[off:off + chunk])
    await ws.send(json.dumps({
        "v": 1, "id": next_id + 1, "type": "media.upload.end",
        "payload": {"media_ref": media_ref, "sha256": hashlib.sha256(data).hexdigest()},
    }))
    print(f"  -> media.upload.end id={next_id + 1} ref={media_ref}")
    while True:
        raw = await asyncio.wait_for(ws.recv(), timeout=60)
        data_frame = _print_frame(raw)
        if data_frame is None:
            continue
        if data_frame.get("type") == "media.upload.ack":
            if not data_frame["payload"].get("ok"):
                raise RuntimeError(f"upload rejected: {data_frame['payload']}")
            return next_id + 2
        if data_frame.get("type") == "error":
            raise RuntimeError(f"upload failed: {data_frame['payload']}")


async def pull_media(ws, media_id: str, request_id: int, expected_size: int | None) -> None:
    """media.pull -> binary frames -> media.pull.end; verifies the size."""
    await ws.send(json.dumps({
        "v": 1, "id": request_id, "type": "media.pull",
        "payload": {"media_id": media_id},
    }))
    print(f"  -> media.pull id={request_id} media_id={media_id}")
    total = 0
    while True:
        raw = await asyncio.wait_for(ws.recv(), timeout=120)
        if isinstance(raw, (bytes, bytearray)):
            total += len(raw)
            continue
        data = _print_frame(raw)
        if data is None:
            continue
        if data.get("type") == "media.pull.end":
            if not data["payload"].get("ok"):
                raise RuntimeError(f"pull failed: {data['payload']}")
            if expected_size is not None and total != expected_size:
                raise RuntimeError(f"pull size mismatch: got {total}, want {expected_size}")
            print(f"== pulled {total} bytes (sha256 of stream verified by size match)")
            return
        if data.get("type") == "error":
            raise RuntimeError(f"pull failed: {data['payload']}")


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},
        },
    }
    if args.fcm_token:
        hello["payload"]["fcm_token"] = args.fcm_token
    await ws.send(json.dumps(hello))
    print("  -> hello" + (f" fcm_token={args.fcm_token[:12]}…" if args.fcm_token else ""))

    # 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

    # M5: optional fcm.register after pairing.
    if args.fcm_reg:
        reg_token = args.fcm_token or f"probe-{uuid.uuid4().hex[:12]}"
        await ws.send(json.dumps({
            "v": 1, "type": "fcm.register",
            "payload": {"fcm_token": reg_token},
        }))
        print(f"  -> fcm.register fcm_token={reg_token[:12]}…")

    # M5: sync catch-up mode (no turn driven).
    if args.sync is not None:
        await ws.send(json.dumps({
            "v": 1, "id": 1, "type": "sync", "payload": {"cursor": args.sync},
        }))
        print(f"  -> sync cursor={args.sync}")
        while True:
            raw = await asyncio.wait_for(ws.recv(), timeout=30)
            data = _print_frame(raw)
            if data is None:
                continue
            if data.get("type") == "sync.done":
                print(f"== sync done at cursor {data['payload'].get('cursor')}")
                await ws.close()
                return 0
            if data.get("type") == "error":
                print(f"!! sync failed: {data['payload']}")
                await ws.close()
                return 8

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

    # M4: optional inbound upload before the turn.
    media_refs: list[str] = []
    next_id = 1
    if args.upload:
        media_ref = f"mu_probe_{uuid.uuid4().hex[:8]}"
        try:
            next_id = await upload_file(ws, args.upload, media_ref, next_id)
        except Exception as e:
            print(f"!! upload failed: {e}")
            await ws.close()
            return 8
        media_refs.append(media_ref)

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

    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 {}
        # M4: fetch offered media live (outbound direction).
        if ftype == "media.offer" and args.pull_offer and payload.get("media_id"):
            try:
                next_id = await pull_media(
                    ws, payload.get("media_id"), next_id, payload.get("size")
                )
            except Exception as e:
                print(f"!! pull failed: {e}")
                await ws.close()
                return 9
        # 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("--upload", default="",
                   help="M4: file to upload (chunked) and attach via media_refs")
    p.add_argument("--pull-offer", action="store_true",
                   help="M4: pull any media.offer that arrives during the turn")
    p.add_argument("--sync", type=int, default=None,
                   help="M5: send sync {cursor} after pairing, print replay, exit")
    p.add_argument("--fcm-token", default="",
                   help="M5: FCM token to attach to the hello payload")
    p.add_argument("--fcm-reg", action="store_true",
                   help="M5: send fcm.register after pairing (uses --fcm-token)")
    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())