#!/usr/bin/env python3
"""Set the RS485 node ID of a ZLAC8015D driver.

Only ONE driver should be connected to the RS485 bus while running this
script (both drivers ship with default ID 1, so with two on the bus their
replies would collide). Recommended workflow to end up with IDs 1 and 2:

    1. Connect ONLY the first driver. Leave it at the default ID:
         python3 zltech_set_id.py --port <PORT> --old-id 1 --new-id 1
       (or just skip it, it is already 1)
    2. Power off, disconnect it, connect ONLY the second driver.
    3. Set its ID to 2:
         python3 zltech_set_id.py --port <PORT> --old-id 1 --new-id 2
    4. Power-cycle that driver.
    5. Now connect both drivers to the same bus; they will answer as
       ID 1 and ID 2 respectively.

Requires: pip install pymodbus pyserial
"""
import argparse
import sys
import time

from zltech_modbus import REG_NODE_ID, REG_SAVE_EEPROM, make_client, read_register, write_register


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--port", required=True, help="Serial port, e.g. /dev/tty.usbserial-XXXX or COM5")
    parser.add_argument("--baud", type=int, default=115200, help="Baud rate (default 115200, the driver default)")
    parser.add_argument("--old-id", type=int, required=True, help="Current node ID of the driver (1-127)")
    parser.add_argument("--new-id", type=int, required=True, help="New node ID to assign (1-127)")
    parser.add_argument("--timeout", type=float, default=0.3, help="Response timeout in seconds")
    args = parser.parse_args()

    for name, value in (("--old-id", args.old_id), ("--new-id", args.new_id)):
        if not (1 <= value <= 127):
            parser.error(f"{name} must be in range 1-127")

    client = make_client(args.port, args.baud, args.timeout)
    if not client.connect():
        print(f"ERROR: could not open serial port {args.port}")
        sys.exit(1)

    try:
        print(f"Checking driver at ID {args.old_id} ...")
        values = read_register(client, args.old_id, REG_NODE_ID, 1)
        if values is None:
            print(f"ERROR: no response from ID {args.old_id}. Check wiring/baud rate, "
                  f"and that exactly one driver is on the bus.")
            sys.exit(1)
        print(f"  driver responded, current node-id register = {values[0]}")

        print(f"Writing new node ID {args.new_id} ...")
        if not write_register(client, args.old_id, REG_NODE_ID, args.new_id):
            print("ERROR: write did not get a valid response.")
            sys.exit(1)

        # The ID change can take effect immediately, so try to save under
        # the new ID first, and fall back to the old ID if that fails.
        time.sleep(0.05)
        saved = write_register(client, args.new_id, REG_SAVE_EEPROM, 1)
        if not saved:
            saved = write_register(client, args.old_id, REG_SAVE_EEPROM, 1)
        if saved:
            print("Saved to EEPROM.")
        else:
            print("WARNING: could not confirm the save-to-EEPROM write; "
                  "the ID change may be lost on power-cycle.")

        print(f"Done. Power-cycle the driver, then verify with:\n"
              f"  python3 zltech_scan_ids.py --port {args.port} --range {args.new_id} {args.new_id}")
    finally:
        client.close()


if __name__ == "__main__":
    main()