#!/usr/bin/env python3
"""Host probe v2 — transport framing variants via the library.

Envelope bytes come from ``EnvelopeCodec``; only the on-the-wire framing
varies (raw, length prefixes, split send, MQTT poke, greeting-first read).
Raw payloads go through ``LocalDeviceClient.send_raw``.

Run from the repo root on the SoftAP-joined host:
    python3 scripts/host_probe2.py --sn A03X2516013210060005 --ports 9000
"""

import argparse
import json
import os
import socket
import sys
import time

sys.path.insert(0, os.path.dirname(os.path.dirname(
    os.path.abspath(__file__))))

from franklinwh_agate_client.device_local import (  # noqa: E402
    LocalDeviceClient,
)
from franklinwh_agate_client.protocol import (  # noqa: E402
    CommunicationCmd,
    direct_connect_request,
    generate_secret_key,
)
from franklinwh_agate_client.protocol.envelope import (  # noqa: E402
    SequenceNumber,
)
from franklinwh_agate_client.protocol.messages import (  # noqa: E402
    EnvelopeCodec,
)

#: Minimal MQTT v3.1.1 CONNECT (client id ``probe01``).
_MQTT_CONNECT = bytes.fromhex("101200044d5154540402003c000670726f62653031")

_seq = SequenceNumber()


def canonical_tx(sn: str, codec: EnvelopeCodec) -> bytes:
    """App-shaped apower query bytes (fresh snno/timeStamp)."""
    env = direct_connect_request(
        CommunicationCmd.DEVICE_APOWER_INFO, sn,
        data={"update": 2}, snno=_seq.next())
    return codec.encode(env)


def recv_all_verbose(sock: socket.socket, idle: float = 4.0):
    out = bytearray()
    sock.settimeout(idle)
    while True:
        try:
            chunk = sock.recv(65536)
        except (socket.timeout, TimeoutError):
            return bytes(out), "idle-timeout"
        if not chunk:
            return bytes(out), "closed-by-peer"
        out.extend(chunk)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--sn", required=True)
    ap.add_argument("--host", default="10.100.1.1")
    ap.add_argument("--ports", default="9000")
    ap.add_argument("--timeout", type=float, default=5.0)
    args = ap.parse_args()
    codec = EnvelopeCodec.for_serial(args.sn)
    print("secret key for %s: %s (radix 16)"
          % (args.sn, generate_secret_key(args.sn)))
    for port in [int(p) for p in args.ports.split(",") if p.strip()]:
        client = LocalDeviceClient(args.sn, host=args.host,
                                   socket_port=port, timeout=args.timeout)
        tx = canonical_tx(args.sn, codec)

        def show(label, tx_bytes, rx, how=""):
            print("--- %s tx=%dB rx=%dB %s"
                  % (label, len(tx_bytes), len(rx), how))
            if rx:
                print("  rx: %s" % rx.hex()[:400])

        # greeting-first: connect, read before sending, then send
        try:
            with socket.create_connection(
                    (args.host, port), timeout=args.timeout) as sock:
                pre, _ = recv_all_verbose(sock, idle=4.0)
                sock.sendall(tx)
                rx, how = recv_all_verbose(sock, idle=8.0)
                show("greet-first/read-then-raw", tx, rx,
                     "pre=%dB %s" % (len(pre), how))
        except (ConnectionError, socket.error, TimeoutError) as exc:
            print("--- greet-first: error: %s" % exc)

        for label, payload in (
            ("raw/long-read", tx),
            ("le16-prefix", len(tx).to_bytes(2, "little") + tx),
            ("be16-prefix", len(tx).to_bytes(2, "big") + tx),
            ("le32-prefix", len(tx).to_bytes(4, "little") + tx),
            ("mqtt-connect", _MQTT_CONNECT),
        ):
            try:
                t0 = time.time()
                rx = client.send_raw(payload, timeout=args.timeout)
                show(label, payload, rx, "%.1fs" % (time.time() - t0))
            except (ConnectionError, socket.error, TimeoutError) as exc:
                print("--- %s: error: %s" % (label, exc))

        # split send over one connection
        try:
            with socket.create_connection(
                    (args.host, port), timeout=args.timeout) as sock:
                half = len(tx) // 2
                sock.sendall(tx[:half])
                time.sleep(0.5)
                sock.sendall(tx[half:])
                rx, how = recv_all_verbose(sock)
                show("raw/split-send", tx, rx, how)
        except (ConnectionError, socket.error, TimeoutError) as exc:
            print("--- split-send: error: %s" % exc)


if __name__ == "__main__":
    main()