"""Battery charge-state proof of concept.

Direct-connect (SoftAP) read path recovered from the app's Dart AOT layer
(see ``docs/RE_DART_FINDINGS.md``): join the device hotspot, open the local
TCP socket, send the encrypted ``CommunicationObject`` envelope, decrypt the
reply, and pull the aggregate state-of-charge out of it.

No cloud, no MQTT, no Bluetooth needed. Requires L2 reachability to the
aGate gateway (``10.100.1.1``) — i.e. run on the host whose Wi-Fi adapter is
associated to the device AP, or in a container bridged onto ``10.100.1.0/24``.
"""

from __future__ import annotations

import argparse
import json
import logging
import socket
from typing import Any, Dict, List, Optional, Tuple

from .device_local import LocalDeviceClient
from .protocol import CommunicationCmd
from .wifi import derive_hotspot_credentials

log = logging.getLogger(__name__)

#: Ports probed in order. ``9000`` was observed open on one live unit
#: (``docs/RE_DART_FINDINGS.md`` §8); ``18000`` (``0x4650``) is the constant
#: in the app's connect path but was closed on that unit.
SOCKET_PORT_CANDIDATES: Tuple[int, ...] = (9000, 18000)

#: Reply keys searched for the aggregate percentage, in priority order.
#: ``batSoc`` is the per-pack charge from the BMS socket query
#: (``queryBmsInfoBySocket`` passes ``arguments={"id": pack}`` — without
#: the pack id the device answers the all-zero ``id: 0`` slot).
SOC_KEYS: Tuple[str, ...] = ("currentBatterySoc", "soc", "batSoc",
                             "batterySoc", "vppSoc")


def _as_dict(data: Any) -> Optional[Dict[str, Any]]:
    if isinstance(data, dict):
        return data
    if isinstance(data, str):
        try:
            parsed = json.loads(data)
        except (ValueError, TypeError):
            return None
        return parsed if isinstance(parsed, dict) else None
    return None


def _walk(obj: Any, key: str) -> List[Any]:
    """Collect every value stored under ``key`` anywhere in nested dicts/lists."""
    found: List[Any] = []
    if isinstance(obj, dict):
        for k, v in obj.items():
            if k == key:
                found.append(v)
            found.extend(_walk(v, key))
    elif isinstance(obj, list):
        for item in obj:
            found.extend(_walk(item, key))
    return found


def _to_float(value: Any) -> Optional[float]:
    try:
        if isinstance(value, bool):
            return None
        num = float(value)
    except (TypeError, ValueError):
        return None
    return num


def extract_soc(reply: Dict[str, Any]) -> Dict[str, Any]:
    """Pull the aggregate charge percentage out of a parsed reply envelope.

    Returns ``{"value": float|None, "source": str|None, "detail": {...}}``.
    ``source`` names the key the value came from; ``detail`` carries the
    supporting fields (``socArray``, ``batteryCount``, charge/discharge
    power) when present.
    """
    data = _as_dict(reply.get("data", reply)) or {}

    for key in SOC_KEYS:
        for container in (data, reply):
            for hit in _walk(container, key):
                num = _to_float(hit)
                if num is not None:
                    return {
                        "value": num,
                        "source": key,
                        "detail": _soc_detail(data),
                    }
    soc_arrays = _walk(data, "socArray")
    for arr in soc_arrays:
        if isinstance(arr, list) and arr:
            nums = [_to_float(v) for v in arr]
            nums = [n for n in nums if n is not None]
            if nums:
                return {
                    "value": sum(nums) / len(nums),
                    "source": "socArray(avg)",
                    "detail": _soc_detail(data),
                }
    return {"value": None, "source": None, "detail": _soc_detail(data)}


def _soc_detail(data: Dict[str, Any]) -> Dict[str, Any]:
    detail: Dict[str, Any] = {}
    for key in (
        "socArray",
        "batteryCount",
        "batteryList",
        "chargePower",
        "dischargePower",
        "currentBatterySoc",
        "soc",
    ):
        hits = _walk(data, key)
        if hits:
            detail[key] = hits[0] if len(hits) == 1 else hits
    # Always record where we looked so failures are debuggable.
    return detail


def get_battery(
    sn: str,
    host: str = "10.100.1.1",
    ports: Tuple[int, ...] = SOCKET_PORT_CANDIDATES,
    key_radix: int = 16,
    timeout: float = 5.0,
) -> Dict[str, Any]:
    """Query the aggregate charge percentage over the local TCP socket.

    Mirrors the app's ``queryAPowerData`` first (``DEVICE_APOWER_INFO`` with
    ``data={"opt": 0, "update": 2}``), then per-pack ``BMS`` socket queries
    (``queryBmsInfoBySocket`` passes the pack ``id`` — ``batSoc`` is the
    pack charge), then ``EMS`` as a fallback, probing each candidate port.
    Like the app's direct-connect flow, a discovery query
    (``QUERY_DEVICE_INFO``, empty serial) opens each port's sequence — its
    reply also confirms two-way crypto before the battery read. Returns
    the ``extract_soc`` dict plus ``cmd``/``port``/``reply`` context.
    Raises the last socket error if no port answers.
    """
    last_err: Optional[Exception] = None
    for port in ports:
        client = LocalDeviceClient(
            sn=sn, host=host, socket_port=port,
            key_radix=key_radix, timeout=timeout,
        )
        try:
            client.send_rpc(CommunicationCmd.QUERY_DEVICE_INFO,
                            equip_no="",
                            data={"minProtocolVer": "V1.00.00"})
        except Exception as exc:  # discovery is best-effort; keep going
            log.debug("port %d discovery failed: %s", port, exc)
        attempts: Tuple[Tuple[CommunicationCmd, Any], ...] = (
            (CommunicationCmd.DEVICE_APOWER_INFO, {"update": 2}),
            (CommunicationCmd.BMS, {"id": 1}),
            (CommunicationCmd.BMS, {"id": 2}),
            (CommunicationCmd.EMS, None),
        )
        for cmd, data in attempts:
            try:
                reply = client.send_rpc(cmd, data=data)
            except (ConnectionError, socket.error, TimeoutError) as exc:
                last_err = exc
                log.debug("port %d cmd %s failed: %s", port, cmd, exc)
                break  # port unreachable — try next port
            except Exception as exc:  # envelope-level failure; try next cmd
                last_err = exc
                log.debug("port %d cmd %s error: %s", port, cmd, exc)
                continue
            result = extract_soc(reply)
            result.update({"cmd": str(cmd), "port": port, "reply": reply})
            if result["value"] is not None:
                return result
            last_err = ValueError(f"reply on port {port} cmd {cmd} has no SoC field")
    raise ConnectionError(f"no battery reply via {host}:{list(ports)}") from last_err


#: ``--probe`` request variants: (label, cmd, data).  The first mirrors the
#: app's ``queryAPowerData`` exactly; the rest isolate single envelope
#: decisions (payload present/empty/absent, fallback commands).
PROBE_VARIANTS: Tuple[Tuple[str, CommunicationCmd, Any], ...] = (
    ("apower/update2", CommunicationCmd.DEVICE_APOWER_INFO, {"update": 2}),
    ("apower/none", CommunicationCmd.DEVICE_APOWER_INFO, None),
    ("apower/empty", CommunicationCmd.DEVICE_APOWER_INFO, {}),
    ("queryDeviceInfo/none", CommunicationCmd.QUERY_DEVICE_INFO, None),
    ("bms/id1", CommunicationCmd.BMS, {"id": 1}),
    ("bms/id2", CommunicationCmd.BMS, {"id": 2}),
    ("ems/none", CommunicationCmd.EMS, None),
)


def probe_device(
    sn: str,
    host: str = "10.100.1.1",
    ports: Tuple[int, ...] = SOCKET_PORT_CANDIDATES,
    key_radix: int = 16,
    timeout: float = 5.0,
) -> List[Dict[str, Any]]:
    """Send every :data:`PROBE_VARIANTS` envelope; return exact TX/RX bytes."""
    results: List[Dict[str, Any]] = []
    for port in ports:
        client = LocalDeviceClient(
            sn=sn, host=host, socket_port=port,
            key_radix=key_radix, timeout=timeout,
        )
        for label, cmd, data in PROBE_VARIANTS:
            try:
                out = client.send_probe(cmd, data=data, timeout=timeout)
            except (ConnectionError, socket.error, TimeoutError) as exc:
                out = {"cmd": str(cmd), "error": f"{type(exc).__name__}: {exc}",
                       "tx": b"", "reply_raw": b"", "reply": None,
                       "envelope": None}
            out.update({"variant": label, "port": port})
            results.append(out)
            if out.get("reply_raw"):
                break  # device answered — deeper variants can wait
    return results


def _recv_all(sock: socket.socket, idle: float = 4.0) -> bytes:
    """Drain a socket until idle timeout or close."""
    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 _apower_tx(sn: str, key_radix: int) -> Tuple[bytes, Dict[str, Any]]:
    """Canonical ``queryAPowerData`` wire bytes (fresh snno/timeStamp)."""
    from .protocol import (
        direct_connect_request,
        encrypt_object,
        generate_secret_key,
    )
    from .protocol.envelope import SequenceNumber

    _apower_tx_seq = getattr(_apower_tx, "_seq", None)
    if _apower_tx_seq is None:
        _apower_tx_seq = _apower_tx._seq = SequenceNumber()  # type: ignore[attr-defined]
    env = direct_connect_request(
        CommunicationCmd.DEVICE_APOWER_INFO, sn,
        data={"update": 2}, snno=_apower_tx_seq.next(),
    )
    wire = encrypt_object(
        env.to_json_string(), generate_secret_key(sn), radix=key_radix)
    return wire, env.to_dict()


#: Minimal MQTT v3.1.1 CONNECT (client id ``probe01``) — rules out / confirms
#: a local broker on the probed port (CONNACK starts ``20 02``).
_MQTT_CONNECT = bytes.fromhex("101200044d5154540402003c000670726f62653031")


def probe_transport(
    sn: str,
    host: str = "10.100.1.1",
    ports: Tuple[int, ...] = (9000,),
    key_radix: int = 16,
    timeout: float = 5.0,
) -> List[Dict[str, Any]]:
    """Transport-level probes against one service port.

    The envelope bytes are already verified byte-for-byte against the
    disassembly; these variants isolate *framing/handshake* behaviour:
    server-first greeting, immediate raw send with a long read, length
    prefixes (LE16/BE16/LE32), an MQTT CONNECT, and a double-send over one
    persistent connection.  Each entry carries exact ``tx``/``rx`` bytes.
    """
    from .protocol import decrypt_data

    results: List[Dict[str, Any]] = []
    for port in ports:
        tx, envelope = _apower_tx(sn, key_radix)

        def run(label: str, tx_bytes: bytes, pre_read: float = 0.0,
                post_idle: float = 4.0, split_send: bool = False) -> None:
            entry: Dict[str, Any] = {
                "variant": label, "port": port, "tx": tx_bytes,
                "rx_pre": b"", "rx": b"", "reply": None, "error": None,
            }
            try:
                with socket.create_connection(
                        (host, port), timeout=timeout) as sock:
                    if pre_read:
                        entry["rx_pre"] = _recv_all(sock, idle=pre_read)
                    if tx_bytes:
                        if split_send:
                            half = len(tx_bytes) // 2
                            sock.sendall(tx_bytes[:half])
                            sock.settimeout(0.5)
                            try:
                                entry["rx_pre"] += sock.recv(65536)
                            except (socket.timeout, TimeoutError):
                                pass
                            sock.sendall(tx_bytes[half:])
                        else:
                            sock.sendall(tx_bytes)
                    entry["rx"] = _recv_all(sock, idle=post_idle)
            except (ConnectionError, socket.error, TimeoutError) as exc:
                entry["error"] = f"{type(exc).__name__}: {exc}"
            if entry["rx"]:
                decoded = decrypt_data(entry["rx"], radix=key_radix)
                if decoded:
                    try:
                        entry["reply"] = json.loads(decoded)
                    except ValueError:
                        entry["reply"] = {"_ undecodable": decoded}
            results.append(entry)

        run("greet-first/read-then-raw", tx, pre_read=4.0)
        run("raw/long-read", tx, post_idle=8.0)
        run("raw/split-send", tx, split_send=True)
        run("le16-prefix", len(tx).to_bytes(2, "little") + tx)
        run("be16-prefix", len(tx).to_bytes(2, "big") + tx)
        run("le32-prefix", len(tx).to_bytes(4, "little") + tx)
        run("mqtt-connect", _MQTT_CONNECT)
    return results


def main(argv: Optional[List[str]] = None) -> int:
    parser = argparse.ArgumentParser(
        description="Read aGate battery percentage over the device SoftAP socket."
    )
    parser.add_argument("--sn", required=True, help="20-char device serial")
    parser.add_argument("--host", default="10.100.1.1", help="aGate gateway IP")
    parser.add_argument(
        "--ports", default=",".join(map(str, SOCKET_PORT_CANDIDATES)),
        help="comma-separated socket ports to probe",
    )
    parser.add_argument("--key-radix", type=int, default=16, choices=(16, 32))
    parser.add_argument("--timeout", type=float, default=5.0)
    parser.add_argument("--show-hotspot", action="store_true",
                        help="print derived AP SSID/password and exit")
    parser.add_argument("--raw", action="store_true", help="print full reply JSON")
    parser.add_argument("--probe", action="store_true",
                        help="send probe variants and print exact TX/RX bytes")
    parser.add_argument("--probe2", action="store_true",
                        help="transport-level probes (greeting, prefixes, "
                             "MQTT, split send) with exact TX/RX bytes")
    args = parser.parse_args(argv)

    if args.show_hotspot:
        profile = derive_hotspot_credentials(args.sn)
        print(json.dumps(
            {"ssid": profile.ssid, "password": profile.password,
             "gateway_ip": profile.gateway_ip}, indent=2))
        return 0

    ports = tuple(int(p) for p in args.ports.split(",") if p.strip())
    if args.probe:
        from .protocol import generate_secret_key
        print(f"secret key for {args.sn}: {generate_secret_key(args.sn)} "
              f"(radix {args.key_radix})")
        for out in probe_device(args.sn, host=args.host, ports=ports,
                                key_radix=args.key_radix, timeout=args.timeout):
            print(f"--- {out['variant']} -> {args.host}:{out['port']}"
                  f" cmd={out['cmd']}")
            if out.get("envelope") is not None:
                print("  envelope: " + json.dumps(out["envelope"], default=str))
            tx = out.get("tx") or b""
            print(f"  tx ({len(tx)} B): {tx.hex()}")
            rx = out.get("reply_raw") or b""
            print(f"  rx ({len(rx)} B): {rx.hex() or '(empty)'}")
            if out.get("reply") is not None:
                print("  reply: " + json.dumps(out["reply"], default=str))
                soc = extract_soc(out["reply"])
                print(f"  soc: {soc['value']} (source={soc['source']})")
            if out.get("error"):
                print(f"  error: {out['error']}")
        return 0
    if args.probe2:
        for out in probe_transport(args.sn, host=args.host, ports=ports,
                                   key_radix=args.key_radix,
                                   timeout=args.timeout):
            print(f"--- {out['variant']} -> {args.host}:{out['port']}")
            print(f"  tx ({len(out['tx'])} B): {out['tx'].hex() or '(none)'}")
            pre = out.get("rx_pre") or b""
            if pre:
                print(f"  rx_pre ({len(pre)} B): {pre.hex()}")
            rx = out.get("rx") or b""
            print(f"  rx ({len(rx)} B): {rx.hex() or '(empty)'}")
            if out.get("reply") is not None:
                print("  reply: " + json.dumps(out["reply"], default=str))
            if out.get("error"):
                print(f"  error: {out['error']}")
        return 0
    try:
        result = get_battery(args.sn, host=args.host, ports=ports,
                             key_radix=args.key_radix, timeout=args.timeout)
    except Exception as exc:
        print(f"ERROR: {exc}")
        return 1
    if args.raw:
        print(json.dumps(result["reply"], indent=2, default=str))
    else:
        value = result["value"]
        print(f"{value:g}%" if value is not None else "unknown")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())