#!/usr/bin/env python3
"""
UPS Modbus RTU 控制执行脚本（用户二次确认后调用）
接收完整的 Modbus RTU 报文字符串，校验后下发给设备。
"""
import argparse
import json
import time
import sys


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


def verify_crc(frame_bytes: bytes) -> bool:
    """验证报文 CRC16 是否正确（最后两字节为 CRC 低字节+高字节）"""
    if len(frame_bytes) < 4:
        return False
    data = frame_bytes[:-2]
    crc = crc16(data)
    crc_lo = crc & 0xFF
    crc_hi = (crc >> 8) & 0xFF
    return frame_bytes[-2] == crc_lo and frame_bytes[-1] == crc_hi


# 白名单：只允许已知的功能码+地址+值组合
WHITELIST = {
    (0x05, 5900, 0xFF00),  # 开机
    (0x05, 5900, 0x0000),  # 停机
    (0x05, 5901, 0xFF00),  # 强制均充
    (0x05, 5902, 0xFF00),  # 开启ECO
    (0x05, 5902, 0x0000),  # 关闭ECO
    (0x06, 5402, 0x0001),  # 开启蜂鸣器
    (0x06, 5402, 0x0002),  # 静音蜂鸣器
    (0x06, 5403, 0x0001),  # 标准电池测试
    (0x06, 5403, 0x0002),  # 深度电池测试
    (0x06, 5403, 0x0003),  # 取消电池测试
    (0x06, 5404, 0x0000),  # CC充电
    (0x06, 5404, 0x0001),  # CP充电
}


def main():
    parser = argparse.ArgumentParser(description="执行 UPS Modbus RTU 控制指令")
    parser.add_argument("--device_id", required=True, help="设备ID或名称")
    parser.add_argument("--command", required=True, help="完整 Modbus RTU 报文，如：01 05 17 0C 00 00 08 7D")
    parser.add_argument("--label", required=True, help="操作说明，如：停机")
    args = parser.parse_args()

    # 解析报文
    try:
        frame_bytes = bytes(int(b, 16) for b in args.command.strip().split())
    except ValueError as e:
        print(json.dumps({"status": "error", "message": f"报文格式错误: {e}"}, ensure_ascii=False))
        sys.exit(1)

    if len(frame_bytes) != 8:
        print(json.dumps({"status": "error", "message": f"报文长度应为8字节，实际{len(frame_bytes)}字节"}, ensure_ascii=False))
        sys.exit(1)

    # CRC 校验
    if not verify_crc(frame_bytes):
        print(json.dumps({"status": "error", "message": "CRC 校验失败，报文可能被篡改"}, ensure_ascii=False))
        sys.exit(1)

    # 白名单校验
    func = frame_bytes[1]
    addr = (frame_bytes[2] << 8) | frame_bytes[3]
    val  = (frame_bytes[4] << 8) | frame_bytes[5]

    if (func, addr, val) not in WHITELIST:
        print(json.dumps({
            "status": "error",
            "message": f"非法指令：func=0x{func:02X} addr={addr} val=0x{val:04X} 不在白名单中"
        }, ensure_ascii=False))
        sys.exit(1)

    # 实际下发（预留接入点，根据网关类型接入 pymodbus 或 HTTP 转发）
    # 示例：通过 Modbus TCP 下发
    # from pymodbus.client import ModbusTcpClient
    # client = ModbusTcpClient(host=DEVICE_IP, port=502)
    # client.connect()
    # if func == 0x05:
    #     client.write_coil(addr, val == 0xFF00, slave=frame_bytes[0])
    # elif func == 0x06:
    #     client.write_register(addr, val, slave=frame_bytes[0])
    # client.close()

    time.sleep(0.3)

    result = {
        "status": "success",
        "message": f"指令下发成功: {args.label}",
        "device_id": args.device_id,
        "modbus_rtu": args.command,
        "parsed": {
            "slave_id": f"0x{frame_bytes[0]:02X}",
            "func": f"0x{func:02X}",
            "addr": addr,
            "val": f"0x{val:04X}",
            "label": args.label
        },
        "timestamp": int(time.time())
    }
    print(json.dumps(result, ensure_ascii=False))


if __name__ == "__main__":
    main()
