#!/usr/bin/env python3
"""PLSR P12 Modbus RTU live-frequency test.

The firmware starts a long Q0 move from S0=D1000 and S1=D1100.  This tool
updates the current-segment frequency at D1010/D1011 with function 0x10, so
the two 16-bit words are committed as one Modbus transaction.
"""

from __future__ import annotations

import argparse
import struct
import time
from dataclasses import dataclass

import serial
from serial.tools import list_ports


S0_BASE = 1000
S1_BASE = 1100
LIVE_FREQUENCY_ADDRESS = S0_BASE + 10
SLAVE_DEFAULT = 1


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


def add_crc(payload: bytes) -> bytes:
    crc = crc16(payload)
    return payload + bytes((crc & 0xFF, crc >> 8))


def signed_dword_words(value: int) -> list[int]:
    raw = value & 0xFFFFFFFF
    return [raw & 0xFFFF, (raw >> 16) & 0xFFFF]


@dataclass
class RtuClient:
    port: serial.Serial
    slave: int

    def exchange(self, request_pdu: bytes, response_size: int) -> bytes:
        request = add_crc(bytes((self.slave,)) + request_pdu)
        self.port.reset_input_buffer()
        self.port.write(request)
        self.port.flush()
        response = self.port.read(response_size)
        if len(response) != response_size:
            raise RuntimeError(
                f"响应超时：期望 {response_size} 字节，收到 {len(response)} 字节"
            )
        if crc16(response[:-2]) != int.from_bytes(response[-2:], "little"):
            raise RuntimeError(f"响应 CRC 错误：{response.hex(' ')}")
        if response[0] != self.slave:
            raise RuntimeError(f"站号错误：收到 {response[0]}，期望 {self.slave}")
        if response[1] & 0x80:
            raise RuntimeError(
                f"Modbus 异常：功能码 0x{response[1]:02X}，异常码 0x{response[2]:02X}"
            )
        return response

    def read_holding(self, address: int, quantity: int) -> list[int]:
        pdu = bytes((0x03,)) + struct.pack(">HH", address, quantity)
        response = self.exchange(pdu, 5 + quantity * 2)
        if response[1] != 0x03 or response[2] != quantity * 2:
            raise RuntimeError(f"0x03 响应格式错误：{response.hex(' ')}")
        return list(struct.unpack(f">{quantity}H", response[3:-2]))

    def write_multiple(self, address: int, values: list[int]) -> None:
        encoded = struct.pack(f">{len(values)}H", *values)
        pdu = (
            bytes((0x10,))
            + struct.pack(">HHB", address, len(values), len(encoded))
            + encoded
        )
        response = self.exchange(pdu, 8)
        expected = bytes((self.slave, 0x10)) + struct.pack(">HH", address, len(values))
        if response[:6] != expected:
            raise RuntimeError(f"0x10 响应回显错误：{response.hex(' ')}")

    def write_dword(self, address: int, value: int) -> None:
        self.write_multiple(address, signed_dword_words(value))
        words = self.read_holding(address, 2)
        if words != signed_dword_words(value):
            raise RuntimeError(
                f"D{address} 回读不一致：写入 {signed_dword_words(value)}，回读 {words}"
            )


def choose_port(requested: str | None) -> str:
    if requested:
        return requested
    ports = [item.device for item in list_ports.comports()]
    if len(ports) == 1:
        print(f"自动选择串口 {ports[0]}")
        return ports[0]
    available = ", ".join(ports) if ports else "未发现串口"
    raise RuntimeError(f"请用 --port 指定串口。当前串口：{available}")


def wait_until(deadline: float) -> None:
    remaining = deadline - time.monotonic()
    if remaining > 0:
        time.sleep(remaining)


def main() -> int:
    parser = argparse.ArgumentParser(description="PLSR P12 Modbus 动态频率自动测试")
    parser.add_argument("--port", help="串口，例如 COM5；只有一个串口时可省略")
    parser.add_argument("--baud", type=int, default=9600)
    parser.add_argument("--slave", type=int, default=SLAVE_DEFAULT)
    args = parser.parse_args()

    port_name = choose_port(args.port)
    with serial.Serial(
        port=port_name,
        baudrate=args.baud,
        bytesize=serial.EIGHTBITS,
        parity=serial.PARITY_EVEN,
        stopbits=serial.STOPBITS_ONE,
        timeout=1.0,
        write_timeout=1.0,
    ) as uart:
        client = RtuClient(uart, args.slave)
        header = client.read_holding(S0_BASE, 20)
        s1 = client.read_holding(S1_BASE, 4)
        if header[0] != 1 or header[12:14] != signed_dword_words(100000):
            raise RuntimeError(
                "P12 数据未就绪：请确认已烧录当前固件并复位开发板"
            )
        if any(s1):
            raise RuntimeError(f"S1 数据异常：{s1}")

        print("P12 已就绪：S0=D1000，S1=D1100，Q0 正在输出")
        print("所有 32 位频率均使用 0x10 一次写入两个寄存器。")
        schedule = [
            (0.0, 1000, "初始目标"),
            (1.0, 4000, "升至 4000Hz"),
            (1.5, 500, "降至 500Hz"),
            (2.0, 0, "0 使用 S2 默认 1000Hz"),
            (2.2, 8000, "超过上限，固件应钳位到 5000Hz"),
            (2.7, -1, "非法值，固件应保持上一次安全目标"),
            (2.9, 2000, "恢复到 2000Hz"),
        ]
        started = time.monotonic()
        for offset, frequency, description in schedule:
            wait_until(started + offset)
            before = time.monotonic()
            client.write_dword(LIVE_FREQUENCY_ADDRESS, frequency)
            latency_ms = (time.monotonic() - before) * 1000.0
            print(
                f"T+{time.monotonic() - started:6.3f}s  "
                f"D1010={frequency:6d}  {description}  RTU往返={latency_ms:6.1f}ms"
            )

        print("脚本测试完成。请按测试说明核对 Q0 波形与 IAR 状态变量。")
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except (RuntimeError, serial.SerialException) as error:
        print(f"测试失败：{error}")
        raise SystemExit(1)