"""Check that a Pylontech US-series battery answers on its RS485 port via a transparent
TCP <-> RS485 channel (Waveshare Transfer Protocol "None").

Sends the Pylontech RS485 "get manufacturer info" (CID2 0x51) and "get analog values"
(CID2 0x42) commands and prints the replies. No third-party packages needed.

Usage:  python pylontech_test.py [host] [port] [first_addr] [last_addr]
        defaults: 192.168.1.203 4196 2 2   (try 2 9 to scan a stack)
"""
import socket
import sys

HOST = sys.argv[1] if len(sys.argv) > 1 else "192.168.1.203"
PORT = int(sys.argv[2]) if len(sys.argv) > 2 else 4196
FIRST = int(sys.argv[3]) if len(sys.argv) > 3 else 2
LAST = int(sys.argv[4]) if len(sys.argv) > 4 else FIRST


def checksum(frame_body):
    """Pylontech frame checksum: two's complement of the ASCII sum, mod 65536."""
    return (~sum(frame_body.encode()) + 1) & 0xFFFF


def length_field(info_hex):
    lenid = len(info_hex)
    lchk = (~((lenid & 0xF) + ((lenid >> 4) & 0xF) + ((lenid >> 8) & 0xF)) + 1) & 0xF
    return f"{lchk:X}{lenid:03X}"


def command(addr, cid2, info_hex=""):
    body = f"20{addr:02X}46{cid2:02X}{length_field(info_hex)}{info_hex}"
    return f"~{body}{checksum(body):04X}\r".encode()


def transact(sock, frame):
    sock.sendall(frame)
    reply = b""
    while not reply.endswith(b"\r"):
        chunk = sock.recv(4096)
        if not chunk:
            raise ConnectionError("gateway closed the connection")
        reply += chunk
    return reply.decode(errors="replace").strip()


def printable(hex_str):
    return "".join(chr(b) for b in bytes.fromhex(hex_str) if 32 <= b < 127).strip()


with socket.create_connection((HOST, PORT), timeout=3) as s:
    s.settimeout(3)
    print(f"Connected to {HOST}:{PORT}")
    for addr in range(FIRST, LAST + 1):
        try:
            reply = transact(s, command(addr, 0x51))
        except socket.timeout:
            print(f"addr {addr}: no reply (check A/B swap, baud rate / DIP switch, address)")
            continue
        rtn, info = reply[7:9], reply[13:-4]
        ok = "OK" if rtn == "00" else f"RTN={rtn}"
        print(f"addr {addr}: {ok}  {printable(info)!r}")
        print(f"   raw: {reply}")
        if rtn == "00":
            analog = transact(s, command(addr, 0x42, f"{addr:02X}"))
            print(f"   analog values raw ({len(analog)} chars): {analog[:80]}...")