import sys
from pymodbus.client import ModbusTcpClient

def read_input_register(host, port, address, slave_id):
    client = ModbusTcpClient(host=host, port=int(port), timeout=3)
    if not client.connect():
        print("终端连接失败")
        return
    try:
        result = client.read_input_registers(address=int(address), count=1, slave=int(slave_id))
        if result.isError():
            print("ZBX_NOTSUPPORTED")
        else:
            print(result.registers[0])
    except:
        print("ZBX_NOTSUPPORTED")
    finally:
        client.close()

def read_discrete_input(host, port, address, slave_id):
    client = ModbusTcpClient(host=host, port=int(port), timeout=3)
    if not client.connect():
        print("终端连接失败")
        return
    try:
        result = client.read_discrete_inputs(address=int(address), count=1, slave=int(slave_id))
        if result.isError():
            print("ZBX_NOTSUPPORTED")
        else:
            print(int(result.bits[0]))  # True/False 转换为 1/0
    except:
        print("ZBX_NOTSUPPORTED")
    finally:
        client.close()

if __name__ == '__main__':
    if len(sys.argv) != 6:
        print("Usage: modbus_query.py <type> <host> <port> <address> <slave_id>")
        sys.exit(1)

    query_type, host, port, address, slave_id = sys.argv[1:]

    if query_type == "input_register":
        read_input_register(host, port, address, slave_id)
    elif query_type == "discrete_input":
        read_discrete_input(host, port, address, slave_id)
    else:
        print("ZBX_NOTSUPPORTED")