#!/usr/bin/env python3
"""Host probe v5 — session sequencing via the library.

One persistent connection carrying discovery → directed queries with
app-like sequential ``snno``, plus close-vs-timeout reporting and
high-snno opens. Payload bytes come from ``direct_connect_request`` +
``EnvelopeCodec``; only the connection orchestration lives here.

Run from the repo root on the SoftAP-joined host:
    python3 scripts/host_probe5.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,
)

_seq = SequenceNumber(start=100)


def tx_for(codec: EnvelopeCodec, sn: str, cmd: CommunicationCmd,
           data, equip=None, snno=None) -> bytes:
    env = direct_connect_request(
        cmd, sn if equip is None else equip, data=data,
        snno=_seq.next() if snno is None else snno)
    return codec.encode(env)


def recv_all_verbose(sock: socket.socket, idle: float = 8.0):
    out = bytearray()
    sock.settimeout(idle)
    while True:
        try:
            chunk = sock.recv(65536)
        except (socket.timeout, TimeoutError):
            return bytes(out), "idle-timeout(%.1fs)" % idle
        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)))
    ports = [int(p) for p in args.ports.split(",") if p.strip()]
    for port in ports:
        print("=== discovery-then-directed -> %s:%d (one connection)"
              % (args.host, port))
        try:
            with socket.create_connection(
                    (args.host, port), timeout=args.timeout) as sock:
                for name, cmd, data, equip in (
                    ("disc/1101+empty", CommunicationCmd.QUERY_DEVICE_INFO,
                     {"minProtocolVer": "V1.00.00"}, ""),
                    ("info/1101+SN", CommunicationCmd.QUERY_DEVICE_INFO,
                     {"minProtocolVer": "V1.00.00"}, None),
                    ("apower/1307", CommunicationCmd.DEVICE_APOWER_INFO,
                     {"update": 2}, None),
                    ("bms/1705", CommunicationCmd.BMS, {"id": 1}, None),
                ):
                    payload = tx_for(codec, args.sn, cmd, data, equip)
                    t0 = time.time()
                    sock.sendall(payload)
                    rx, how = recv_all_verbose(sock)
                    print("--- %s tx=%dB rx=%dB (%.1fs, %s)"
                          % (name, len(payload), len(rx),
                             time.time() - t0, how))
                    if rx:
                        print("  rx: %s" % rx.hex()[:400])
                        parsed = codec.try_decode(rx)
                        if parsed is not None:
                            print("  reply: " + json.dumps(
                                parsed, default=str)[:800])
                    if how == "closed-by-peer" and not rx:
                        print("  (peer closed; stopping session)")
                        break
        except (ConnectionError, socket.error, TimeoutError) as exc:
            print("  error: %s: %s" % (type(exc).__name__, exc))

        client = LocalDeviceClient(args.sn, host=args.host,
                                   socket_port=port, timeout=args.timeout)
        for label, cmd, data, snno in (
            ("apower-high-snno", CommunicationCmd.DEVICE_APOWER_INFO,
             {"update": 2}, 500),
            ("info-high-snno", CommunicationCmd.QUERY_DEVICE_INFO,
             {"minProtocolVer": "V1.00.00"}, 501),
        ):
            env = direct_connect_request(cmd, args.sn, data=data, snno=snno)
            payload = codec.encode(env)
            try:
                rx = client.send_raw(payload, timeout=args.timeout)
            except (ConnectionError, socket.error, TimeoutError) as exc:
                print("--- %s: error: %s" % (label, exc))
                continue
            print("--- %s tx=%dB rx=%dB" % (label, len(payload), len(rx)))
            parsed = codec.try_decode(rx)
            if parsed is not None:
                print("  reply: " + json.dumps(parsed, default=str)[:800])


if __name__ == "__main__":
    main()