#!/usr/bin/env python3
"""Bench smoke-test for the modbus_2_io device over Modbus TCP.

Reads the config + I/O registers, exercises an output, and checks status bits.
Discovers the device via mDNS if --host is omitted.

  pip install pymodbus zeroconf
  python3 tools/bench_test.py                      # auto-discover via mDNS
  python3 tools/bench_test.py --host 192.168.1.50  # direct
  python3 tools/bench_test.py --model e04          # E04SD register map

Verifies Phase 7 items: register read/write, output drive, status flags.
RTU testing needs a USB-485 dongle (use pymodbus serial client separately).
"""
import argparse, sys, time

E02_REGS = {
    0: "unit_id", 1: "in1_mode", 2: "in2_mode", 3: "failsafe_out",
    4: "failsafe_tmo", 5: "rtu_enable", 6: "rtu_baud/100", 7: "rtu_parity",
    8: "rtu_stopbits", 9: "in1_value", 10: "in2_value", 11: "output%",
    12: "status_flags",
}
E04_REGS = {
    0: "unit_id", 1: "failsafe_mask", 2: "failsafe_tmo", 3: "rtu_enable",
    4: "rtu_baud/100", 5: "rtu_parity", 6: "rtu_stopbits",
    7: "in1", 8: "in2", 9: "in3", 10: "in4",
    11: "out1", 12: "out2", 13: "out3", 14: "out4", 15: "status_flags",
}
STATUS_BITS = {
    0x01: "LINK_UP", 0x02: "MASTER_ACTIVE", 0x04: "IN1_WIRE_BREAK",
    0x08: "IN2_WIRE_BREAK", 0x10: "FAILSAFE_ACTIVE", 0x20: "RTU_ENABLED",
}


def discover(timeout=5.0):
    try:
        from zeroconf import Zeroconf, ServiceBrowser
    except ImportError:
        sys.exit("zeroconf not installed: pip install zeroconf  (or pass --host)")
    import socket
    found = []

    class L:
        def add_service(self, zc, t, name):
            info = zc.get_service_info(t, name)
            if info and info.addresses:
                ip = socket.inet_ntoa(info.addresses[0])
                txt = {k.decode(): v.decode() for k, v in info.properties.items()}
                found.append((ip, info.port, txt))

        def update_service(self, *a): pass
        def remove_service(self, *a): pass

    zc = Zeroconf()
    ServiceBrowser(zc, "_edge_modbus._tcp.local.", L())
    print(f"Browsing _edge_modbus._tcp for {timeout}s ...")
    time.sleep(timeout)
    zc.close()
    return found


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--host")
    ap.add_argument("--port", type=int, default=502)
    ap.add_argument("--unit", type=int, default=1)
    ap.add_argument("--model", choices=["e02", "e04"], default="e02")
    ap.add_argument("--drive", type=int, help="write this value to output 1 and read back")
    args = ap.parse_args()

    host = args.host
    if not host:
        devs = discover()
        if not devs:
            sys.exit("No device found via mDNS. Pass --host.")
        for ip, port, txt in devs:
            print(f"  found {txt.get('model','?')} @ {ip}:{port}  TXT={txt}")
        host, args.port = devs[0][0], devs[0][1]
        print(f"Using {host}:{args.port}")

    try:
        from pymodbus.client import ModbusTcpClient
    except ImportError:
        sys.exit("pymodbus not installed: pip install pymodbus")

    regs = E04_REGS if args.model == "e04" else E02_REGS
    n = max(regs) + 1
    c = ModbusTcpClient(host, port=args.port, timeout=3)
    if not c.connect():
        sys.exit(f"connect failed to {host}:{args.port}")

    rr = c.read_holding_registers(0, count=n, slave=args.unit)
    if rr.isError():
        sys.exit(f"read error: {rr}")
    print("\nHolding registers:")
    for a in range(n):
        print(f"  [{a:2}] {regs.get(a,'?'):>14} = {rr.registers[a]}")

    status = rr.registers[max(regs)]
    flags = [name for bit, name in STATUS_BITS.items() if status & bit]
    print(f"\nStatus 0x{status:02X}: {', '.join(flags) or '(none)'}")

    if args.drive is not None:
        out_addr = 11
        print(f"\nWriting output[{out_addr}] <- {args.drive}")
        wr = c.write_register(out_addr, args.drive, slave=args.unit)
        if wr.isError():
            sys.exit(f"write error: {wr}")
        time.sleep(0.2)
        rb = c.read_holding_registers(out_addr, count=1, slave=args.unit)
        print(f"read back = {rb.registers[0]}  ({'OK' if rb.registers[0]==args.drive else 'MISMATCH'})")

    c.close()
    print("\ndone.")


if __name__ == "__main__":
    main()