#!/usr/bin/env python3
import argparse, socket, struct, time, json, datetime

MAX_BLOCK = 120  # pysy <125

def crc16_modbus(data: bytes) -> int:
    crc = 0xFFFF
    for b in data:
        crc ^= b
        for _ in range(8):
            if crc & 1:
                crc = (crc >> 1) ^ 0xA001
            else:
                crc >>= 1
    return crc & 0xFFFF

def build_read_req(unit: int, func: int, addr: int, count: int) -> bytes:
    # RTU: [unit][func][addr_hi][addr_lo][count_hi][count_lo][crc_lo][crc_hi]
    pdu = struct.pack(">BBHH", unit, func, addr, count)
    crc = crc16_modbus(pdu)
    return pdu + struct.pack("<H", crc)

def recv_exact(sock: socket.socket, n: int, timeout: float) -> bytes:
    sock.settimeout(timeout)
    out = b""
    while len(out) < n:
        chunk = sock.recv(n - len(out))
        if not chunk:
            raise TimeoutError("socket closed/no data")
        out += chunk
    return out

def read_regs_rtuotcp(sock: socket.socket, unit: int, func: int, addr: int, count: int, timeout=1.5):
    req = build_read_req(unit, func, addr, count)
    sock.sendall(req)

    # Response:
    # normal: [unit][func][bytecount][data...][crc_lo][crc_hi]
    # exception: [unit][func|0x80][excode][crc_lo][crc_hi]
    hdr = recv_exact(sock, 3, timeout)
    ru, rf, bc = hdr[0], hdr[1], hdr[2]

    if ru != unit:
        # joskus linjalla voi olla "jäänteitä" -> nostetaan virhe ja retry
        raise ValueError(f"unit mismatch: got {ru} expected {unit}")

    if rf & 0x80:
        # exception
        rest = recv_exact(sock, 2, timeout)  # excode + crc(2)?? (excode=1, crc=2 -> total 3, but hdr had 3 -> bc is actually excode here)
        # korjaus: exception frame on 5 bytes total: unit, func|80, excode, crc_lo, crc_hi
        # meillä hdr=3 sisältää excode bc:ssä
        frame = hdr + rest  # total 5
        crc_recv = struct.unpack("<H", frame[-2:])[0]
        crc_calc = crc16_modbus(frame[:-2])
        if crc_recv != crc_calc:
            raise ValueError("CRC error (exception)")
        raise RuntimeError(f"Modbus exception func=0x{rf:02X} code=0x{bc:02X}")

    # normal: bc = 2*count
    data_len = bc + 2  # data + crc
    rest = recv_exact(sock, data_len, timeout)
    frame = hdr + rest
    crc_recv = struct.unpack("<H", frame[-2:])[0]
    crc_calc = crc16_modbus(frame[:-2])
    if crc_recv != crc_calc:
        raise ValueError("CRC error")

    data = frame[3:-2]
    if len(data) != 2 * count:
        raise ValueError(f"length mismatch: got {len(data)} expected {2*count}")

    regs = list(struct.unpack(">" + "H"*count, data))
    return regs

def dump_range(host, port, unit, kind, start, end, timeout, delay, retries):
    # kind: "input" => func 0x04, "holding" => func 0x03
    func = 0x04 if kind == "input" else 0x03
    out = {}
    addr = start

    with socket.create_connection((host, port), timeout=timeout) as sock:
        while addr <= end:
            n = min(MAX_BLOCK, end - addr + 1)

            last_err = None
            for _ in range(retries + 1):
                try:
                    regs = read_regs_rtuotcp(sock, unit, func, addr, n, timeout=timeout)
                    for i, v in enumerate(regs):
                        out[addr + i] = int(v)
                    last_err = None
                    break
                except Exception as e:
                    last_err = str(e)
                    time.sleep(0.15)

            if last_err is not None:
                out[f"err_{kind}_{addr}_{n}"] = last_err

            addr += n
            time.sleep(delay)

    return out

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--host", default="192.168.1.20")
    ap.add_argument("--port", type=int, default=5020)
    ap.add_argument("--unit", type=int, default=1)

    ap.add_argument("--in_start", type=int, default=0)
    ap.add_argument("--in_end", type=int, default=400)

    ap.add_argument("--hold_start", type=int, default=0)
    ap.add_argument("--hold_end", type=int, default=400)

    ap.add_argument("--timeout", type=float, default=1.8)
    ap.add_argument("--delay", type=float, default=0.04)   # pienempi = nopeampi, isompi = varmempi
    ap.add_argument("--retries", type=int, default=2)
    ap.add_argument("--out", default=None)
    args = ap.parse_args()

    stamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    outpath = args.out or f"./dump_{stamp}.json"

    data = {
        "meta": {
            "ts": stamp,
            "transport": "rtu_over_tcp",
            "host": args.host,
            "port": args.port,
            "unit": args.unit,
            "input_range": [args.in_start, args.in_end],
            "holding_range": [args.hold_start, args.hold_end],
            "timeout": args.timeout,
            "delay": args.delay,
            "retries": args.retries,
        },
        "input": dump_range(args.host, args.port, args.unit, "input", args.in_start, args.in_end, args.timeout, args.delay, args.retries),
        "holding": dump_range(args.host, args.port, args.unit, "holding", args.hold_start, args.hold_end, args.timeout, args.delay, args.retries),
    }

    with open(outpath, "w") as f:
        json.dump(data, f, indent=2, sort_keys=True)

    print("Wrote:", outpath)

if __name__ == "__main__":
    main()