#!/usr/bin/env python3
"""Host probe v6 — verified-against-capture wire format (stdlib only).

Locked by two live frames (tests/fixtures/*.bin in the repo):
* key radix 16 (``int("80",16)=128``, ``int("1c",16)=28``);
* ``S1 || transform(S2)`` both directions; reply id = request id + 1
  observed (1101 -> 1102);
* socket payloads seeded with ``"opt": 0``;
* replies parsed by splitting at the 2nd comma (no ids needed up front).

Run on the SoftAP-joined host:
    python3 host_probe6.py --sn A03X2516013210060005 --ports 9000
"""

import argparse
import json
import socket
import time
import zlib

APOWER, QUERY_INFO, BMS, EMS = 0x51B, 0x44D, 0x6A9, 0x6AB
SNNO = [0]
SOC_KEYS = ("currentBatterySoc", "soc", "batterySoc", "vppSoc")


def secret_key(sn):
    return format(sum(sn.encode()) & 0xFF, "x")


def build_envelope(sn, cmd, args=None, snno=None, ts=None, equip=None):
    if snno is None:
        SNNO[0] = (SNNO[0] % 9999) + 1
        snno = SNNO[0]
    data = {"opt": 0}
    if args:
        data.update(args)
    compact = json.dumps(data, separators=(",", ":"))
    env = {"cmdType": cmd, "equipNo": equip if equip is not None else sn,
           "type": 0,
           "timeStamp": ts if ts is not None else int(time.time()),
           "snno": snno,
           "len": len(compact.encode()),
           "crc": format(zlib.crc32(compact.encode()), "X"),
           "dataArea": data}
    return json.dumps(env, separators=(",", ":"))


def encrypt(json_str, key):
    parts = json_str.split(",")
    a = (parts[0] + "," + parts[1] + ",").encode()
    b = ",".join(parts[2:]).encode()
    n = int(key, 16)
    return a + bytes((b[i] + i + len(a) + n) & 0xFF
                     for i in range(len(b)))


def decrypt_wire(wire):
    first = wire.find(b",")
    second = wire.find(b",", first + 1)
    if first < 0 or second < 0:
        return ""
    a, b = wire[:second + 1], wire[second + 1:]
    try:
        probe = json.loads(a[:-1].decode() + "}")
    except Exception:
        return ""
    if not isinstance(probe, dict):
        return ""
    n = int(secret_key(str(probe.get("equipNo", ""))), 16)
    try:
        return (a + bytes((t - i - len(a) - n) & 0xFF
                          for i, t in enumerate(b))).decode()
    except Exception:
        return ""


def recv_all(sock, idle=8.0):
    out = bytearray()
    sock.settimeout(idle)
    while True:
        try:
            chunk = sock.recv(65536)
        except (socket.timeout, TimeoutError):
            break
        if not chunk:
            break
        out.extend(chunk)
    return bytes(out)


def walk(obj, key, found):
    if isinstance(obj, dict):
        for k, v in obj.items():
            if k == key:
                found.append(v)
            walk(v, key, found)
    elif isinstance(obj, list):
        for item in obj:
            walk(item, key, found)
    return found


def extract_soc(reply):
    data = reply.get("dataArea", reply)
    if isinstance(data, str):
        try:
            data = json.loads(data)
        except ValueError:
            data = {}
    for key in SOC_KEYS:
        hits = walk(reply, key, [])
        for hit in hits:
            if isinstance(hit, bool):
                continue
            try:
                return float(hit), key
            except (TypeError, ValueError):
                continue
    arrs = walk(data if isinstance(data, dict) else {}, "socArray", [])
    for arr in arrs:
        if isinstance(arr, list) and arr:
            try:
                nums = [float(v) for v in arr
                        if not isinstance(v, bool)]
                if nums:
                    return sum(nums) / len(nums), "socArray(avg)"
            except (TypeError, ValueError):
                continue
    return None, None


def run(label, host, port, tx, timeout, pre_read=0.0):
    rx_pre, rx, err, reply = b"", b"", None, None
    try:
        with socket.create_connection((host, port),
                                      timeout=timeout) as sock:
            if pre_read:
                rx_pre = recv_all(sock, idle=pre_read)
            if tx:
                sock.sendall(tx)
                rx = recv_all(sock)
            if rx:
                dec = decrypt_wire(rx)
                reply = json.loads(dec) if dec else {"_raw": rx.hex()}
    except (ConnectionError, socket.error, TimeoutError) as exc:
        err = "%s: %s" % (type(exc).__name__, exc)
    print("--- %s -> %s:%d" % (label, host, port))
    print("  tx (%d B): %s" % (len(tx), tx.hex() or "(none)"))
    if rx_pre:
        print("  rx_pre (%d B): %s" % (len(rx_pre), rx_pre.hex()))
    print("  rx (%d B): %s" % (len(rx), rx.hex()[:400] or "(empty)"))
    if reply is not None:
        print("  reply: " + json.dumps(reply, default=str)[:3000])
        if isinstance(reply, dict):
            val, src = extract_soc(reply)
            print("  soc: %s (source=%s)" % (val, src))
    if err:
        print("  error: %s" % err)
    return rx


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()
    key = secret_key(args.sn)
    print("secret key for %s: %s (radix 16)" % (args.sn, key))
    variants = [
        ("discovery/1101+empty", QUERY_INFO,
         {"minProtocolVer": "V1.00.00"}, ""),
        ("queryInfo/1101+SN", QUERY_INFO,
         {"minProtocolVer": "V1.00.00"}, None),
        ("apower/1307+SN", APOWER, {"update": 2}, None),
        ("apower/1307+IBG", APOWER, {"update": 2},
         "10060005A03X25160132"),
        ("bms/1705+SN", BMS, None, None),
        ("ems/1707+SN", EMS, None, None),
    ]
    for port in [int(p) for p in args.ports.split(",") if p.strip()]:
        for label, cmd, data, equip in variants:
            tx = encrypt(build_envelope(
                args.sn, cmd, data,
                equip=args.sn if equip is None else equip or "00000000",
            ), key)
            if run(label, args.host, port, tx, args.timeout,
                   pre_read=4.0 if label.startswith("discovery") else 0.0):
                pass


if __name__ == "__main__":
    main()