#!/usr/bin/env python3
"""PLSR P14 four-axis 100 kHz hardware-counter stress test."""

from __future__ import annotations

import argparse
import time

import serial

from plsr_modbus_frequency_test import RtuClient, choose_port, signed_dword_words


CONTROL_BASE = 1200
CONTROL_WINDOW_WORDS = 338
CALL_REQUEST = CONTROL_BASE + 8
CALL_RESPONSE = CONTROL_BASE + 24
COMMAND_REQUEST = CONTROL_BASE + 40
COMMAND_RESPONSE = CONTROL_BASE + 48
AXIS_STATUS_BASE = CONTROL_BASE + 64
AXIS_STATUS_WORDS = 48
PERFORMANCE_BASE = CONTROL_BASE + 56
STAGE_PERFORMANCE_BASE = CONTROL_BASE + 256
AB_GATE_PERFORMANCE_BASE = CONTROL_BASE + 268
PERFORMANCE_VERSION = 7
STAGE_NAMES = (
    "脉冲合并/保护",
    "关键事件",
    "命令队列",
    "普通事件/方向提交",
    "HAL/路径/Profile",
    "HSD检查点",
)
STAGE_PERFORMANCE_WORDS = len(STAGE_NAMES) * 2

S0_BASES = (1600, 1800, 2000, 2200)
S1_BASES = (1700, 1900, 2100, 2300)
TEST_FREQUENCY_HZ = 100_000
TEST_PULSES = 200_000

RESULT_OK = 0
RESULT_QUEUED = 1
STATE_ACCEL = 2
STATE_RUN = 3
STATE_COMPLETED = 7

CALL_COMMIT = 1
CALL_START = 2
CMD_SET_POSITION = 5


def put_u32(words: list[int], offset: int, value: int) -> None:
    words[offset : offset + 2] = signed_dword_words(value)


def put_u64(words: list[int], offset: int, value: int) -> None:
    raw = value & 0xFFFFFFFFFFFFFFFF
    words[offset : offset + 4] = [
        (raw >> shift) & 0xFFFF for shift in (0, 16, 32, 48)
    ]


def get_u32(words: list[int], offset: int) -> int:
    return words[offset] | (words[offset + 1] << 16)


def get_u64(words: list[int], offset: int, signed: bool = False) -> int:
    raw = sum(words[offset + index] << (16 * index) for index in range(4))
    if signed and raw & (1 << 63):
        return raw - (1 << 64)
    return raw


def wait_response(
    client: RtuClient, address: int, words: int, sequence: int, timeout: float = 2.0
) -> list[int]:
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        response = client.read_holding(address, words)
        if get_u32(response, 0) == sequence:
            return response
    raise RuntimeError(f"等待序号 {sequence} 的应答超时")


def send_command(
    client: RtuClient, sequence: int, axis: int, opcode: int, argument: int = 0
) -> list[int]:
    request = [0] * 8
    put_u32(request, 0, sequence)
    request[2] = opcode
    request[3] = axis
    put_u64(request, 4, argument)
    client.write_multiple(COMMAND_REQUEST, request)
    return wait_response(client, COMMAND_RESPONSE, 8, sequence)


def send_call(
    client: RtuClient,
    sequence: int,
    axis: int,
    operation: int,
    s2_set: int = 1,
) -> list[int]:
    request = [0] * 16
    put_u32(request, 0, sequence)
    request[2] = 0  # S0 device D
    put_u32(request, 3, S0_BASES[axis])
    request[5] = 0  # S1 device D
    put_u32(request, 6, S1_BASES[axis])
    request[8] = 0  # S2 constant
    put_u32(request, 10, s2_set)
    request[12] = axis
    request[13] = 0  # PULSE/DIR
    request[14] = operation
    client.write_multiple(CALL_REQUEST, request)
    return wait_response(client, CALL_RESPONSE, 12, sequence)


def check_result(response: list[int], offset: int, expected: int, label: str) -> None:
    if response[offset] != expected:
        raise RuntimeError(
            f"{label} 返回 {response[offset]}，期望 {expected}；应答={response}"
        )


def read_axis_status(client: RtuClient, axis: int) -> dict[str, int]:
    address = AXIS_STATUS_BASE + axis * AXIS_STATUS_WORDS
    words = client.read_holding(address, AXIS_STATUS_WORDS)
    generation_begin = get_u32(words, 0)
    generation_end = get_u32(words, 46)
    if generation_begin != generation_end or generation_begin & 1:
        raise RuntimeError(
            f"轴{axis}状态快照不一致：begin={generation_begin}, end={generation_end}"
        )
    return {
        "state": words[2],
        "flags": get_u32(words, 3),
        "error": words[6],
        "stop_reason": words[7],
        "last_result": words[8],
        "last_sequence": get_u32(words, 10),
        "logical_position": get_u64(words, 16, signed=True),
        "task_pulses": get_u64(words, 20, signed=True),
        "physical_pulses": get_u64(words, 28),
        "counter_mode": words[37],
        "current_frequency": get_u32(words, 38),
        "target_frequency": get_u32(words, 40),
    }


def print_process_checkpoint(client: RtuClient, label: str) -> None:
    words = client.read_holding(PERFORMANCE_BASE, 4)
    stage_words = client.read_holding(
        STAGE_PERFORMANCE_BASE, STAGE_PERFORMANCE_WORDS
    )
    stages = ", ".join(
        f"{name}={get_u32(stage_words, index * 2)}"
        for index, name in enumerate(STAGE_NAMES)
    )
    print(
        f"P16阶段[{label}]：自身最大={get_u32(words, 0)} cycles，"
        f"响应最大={get_u32(words, 2)} cycles"
    )
    print(f"  分段最大：{stages}")


def prepare_jobs(client: RtuClient) -> None:
    for axis in range(4):
        s0 = [0] * 20
        put_u32(s0, 0, 1)
        put_u32(s0, 10, TEST_FREQUENCY_HZ)
        put_u32(s0, 12, TEST_PULSES)
        client.write_multiple(S0_BASES[axis], s0)
        client.write_multiple(S1_BASES[axis], [0] * 4)


def main() -> int:
    parser = argparse.ArgumentParser(
        description="PLSR P14 四轴100kHz硬件计数与并发压力测试"
    )
    parser.add_argument("--port", help="串口，例如 COM5；只有一个串口时可省略")
    parser.add_argument("--baud", type=int, default=9600)
    parser.add_argument("--slave", type=int, default=1)
    args = parser.parse_args()

    with serial.Serial(
        port=choose_port(args.port),
        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(CONTROL_BASE, 8)
        if header[:5] != [
            0x504C,
            0x5352,
            0x0100,
            CONTROL_WINDOW_WORDS,
            0x0007,
        ]:
            raise RuntimeError(
                f"P14控制窗口未就绪：{header}；请烧录当前固件并复位"
            )
        print("P14 已就绪：四轴 PULSE/DIR，100kHz，200000脉冲/轴")

        prepare_jobs(client)
        sequence = 100
        for axis in range(4):
            response = send_command(
                client, sequence, axis, CMD_SET_POSITION, argument=0
            )
            check_result(response, 4, RESULT_QUEUED, f"轴{axis} SET_POSITION")
            sequence += 1
        print("四轴位置已清零，S0/S1 已用 0x10 原子写入")
        print_process_checkpoint(client, "位置清零")

        for axis in range(4):
            response = send_call(client, sequence, axis, CALL_COMMIT)
            check_result(response, 3, RESULT_OK, f"轴{axis} COMMIT")
            if response[11] != 1:
                raise RuntimeError(f"轴{axis} COMMIT 未建立有效快照")
            sequence += 1
        print("四轴 COMMIT 校验通过")
        print_process_checkpoint(client, "COMMIT")

        physical_baseline = [
            read_axis_status(client, axis)["physical_pulses"]
            for axis in range(4)
        ]

        started = time.monotonic()
        for axis in range(4):
            response = send_call(client, sequence, axis, CALL_START)
            check_result(response, 3, RESULT_QUEUED, f"轴{axis} START")
            sequence += 1
        print("四轴 START 已排队；持续读取状态以施加 Modbus/任务并发压力")
        print_process_checkpoint(client, "START")

        running_status = [read_axis_status(client, axis) for axis in range(4)]
        for axis, status in enumerate(running_status):
            if status["state"] not in {STATE_ACCEL, STATE_RUN}:
                raise RuntimeError(f"轴{axis} 未进入运行态：{status}")
            expected_mode = 1 if axis < 2 else 0
            if status["counter_mode"] != expected_mode:
                raise RuntimeError(
                    f"轴{axis}计数模式={status['counter_mode']}，期望={expected_mode}"
                )
            if status["current_frequency"] != TEST_FREQUENCY_HZ:
                raise RuntimeError(f"轴{axis}频率不正确：{status}")
        print("运行期计数租约正确：Q0/Q1=硬件，Q2/Q3=软件回退")
        print_process_checkpoint(client, "进入运行态")

        deadline = time.monotonic() + 10.0
        polls = 0
        final_status: list[dict[str, int]] = running_status
        while time.monotonic() < deadline:
            final_status = [read_axis_status(client, axis) for axis in range(4)]
            polls += 4
            if all(item["state"] == STATE_COMPLETED for item in final_status):
                break
        else:
            raise RuntimeError(f"等待四轴完成超时，最后状态：{final_status}")

        for axis, status in enumerate(final_status):
            expected = {
                "logical_position": TEST_PULSES,
                "task_pulses": TEST_PULSES,
                "error": 0,
                "last_result": RESULT_OK,
            }
            bad = {key: (status[key], value) for key, value in expected.items()
                   if status[key] != value}
            physical_delta = (
                status["physical_pulses"] - physical_baseline[axis]
            )
            if physical_delta != TEST_PULSES:
                bad["physical_pulses_delta"] = (
                    physical_delta,
                    TEST_PULSES,
                )
            if bad:
                raise RuntimeError(f"轴{axis}最终计数不正确：{bad}；状态={status}")

        elapsed = time.monotonic() - started
        print_process_checkpoint(client, "运行完成")
        performance = client.read_holding(PERFORMANCE_BASE, 8)
        performance_again = client.read_holding(PERFORMANCE_BASE, 8)
        if performance != performance_again:
            performance = performance_again
        stage_performance = client.read_holding(
            STAGE_PERFORMANCE_BASE, STAGE_PERFORMANCE_WORDS
        )
        ab_gate_cycles = get_u32(
            client.read_holding(AB_GATE_PERFORMANCE_BASE, 2), 0
        )
        core_clock_hz = get_u32(header, 5)
        performance_version = header[7]
        cycle_values = {
            "PlsrProcess自身": get_u32(performance, 0),
            "PlsrProcess响应": get_u32(performance, 2),
            "TIM6控制ISR": get_u32(performance, 4),
            "输出定时器ISR": performance[6],
            "TIM9/12计数ISR": performance[7],
        }
        if core_clock_hz == 0 or performance_version != PERFORMANCE_VERSION:
            raise RuntimeError(
                f"P16性能诊断头无效：clock={core_clock_hz}, "
                f"version={performance_version}"
            )
        if any(value == 0 for value in cycle_values.values()):
            raise RuntimeError(f"P16性能计数未完整运行：{cycle_values}")
        budgets = {
            "PlsrProcess自身": core_clock_hz // 1_000,
            "TIM6控制ISR": core_clock_hz // 10_000,
            "输出定时器ISR": core_clock_hz // TEST_FREQUENCY_HZ,
            "TIM9/12计数ISR": core_clock_hz // TEST_FREQUENCY_HZ,
        }
        overruns = {
            name: (cycles, budgets[name])
            for name, cycles in cycle_values.items()
            if name in budgets
            if cycles >= budgets[name]
        }
        print("P16 DWT最坏执行时间：")
        for name, cycles in cycle_values.items():
            microseconds = cycles * 1_000_000.0 / core_clock_hz
            budget_text = (
                f"预算<{budgets[name]} cycles"
                if name in budgets
                else "观测项（含中断抢占）"
            )
            print(
                f"  {name:<16} {cycles:8d} cycles  "
                f"{microseconds:8.3f}us  {budget_text}"
            )
        print("P16 PlsrProcess分段最大执行时间（各段独立峰值）：")
        for index, name in enumerate(STAGE_NAMES):
            cycles = get_u32(stage_performance, index * 2)
            microseconds = cycles * 1_000_000.0 / core_clock_hz
            print(f"  {name:<18} {cycles:8d} cycles  {microseconds:8.3f}us")
        print(
            "  AB末周期快速门控       "
            f"{ab_gate_cycles:8d} cycles  "
            f"{ab_gate_cycles * 1_000_000.0 / core_clock_hz:8.3f}us  "
            "（PULSE/DIR用例未执行时允许为0）"
        )
        if overruns:
            raise RuntimeError(f"实时执行时间超过对应调度周期：{overruns}")
        response_cycles = cycle_values["PlsrProcess响应"]
        if response_cycles >= core_clock_hz // 1_000:
            print(
                "提示：PlsrProcess墙钟响应超过1ms，但自身CPU执行时间达标；"
                "差值来自高优先级PLSR定时器中断抢占。"
            )
        print(
            f"全部 PASS：四轴均为 {TEST_PULSES} 脉冲，"
            f"耗时 {elapsed:.3f}s，运行期状态读取 {polls} 次"
        )
        print("请再核对逻辑分析仪：Q0~Q3各200000个上升沿、100kHz、无窄脉冲。")
    return 0


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