#!/usr/bin/env python3
"""PLSR dual-AB 100 kHz hardware-counter and fast-gate stress test."""

from __future__ import annotations

import argparse
import time
from dataclasses import dataclass

import serial

from plsr_modbus_counter_stress_test import (
    AB_GATE_PERFORMANCE_BASE,
    AXIS_STATUS_WORDS,
    CALL_COMMIT,
    CALL_REQUEST,
    CALL_RESPONSE,
    CALL_START,
CMD_SET_POSITION,
    CONTROL_BASE,
    CONTROL_WINDOW_WORDS,
    PERFORMANCE_BASE,
    PERFORMANCE_VERSION,
    RESULT_OK,
    RESULT_QUEUED,
    RtuClient,
    S0_BASES,
    S1_BASES,
    check_result,
    get_u32,
    put_u32,
    read_axis_status,
    send_command,
    wait_response,
)
from plsr_modbus_frequency_test import choose_port


TEST_FREQUENCY_HZ = 100_000
TEST_LOW_FREQUENCY_HZ = 50_000
AB_OUTPUT_MODE = 1
AB_S2_SET = 3
AB_OWNERS = (0, 2)
CMD_PAUSE = 3
CMD_RESUME = 4
STATE_ACCEL = 2
STATE_RUN = 3
STATE_DECEL = 4
STATE_PAUSED = 6
STATE_COMPLETED = 7
RUNNING_STATES = {STATE_ACCEL, STATE_RUN, STATE_DECEL}


@dataclass(frozen=True)
class AbCase:
    name: str
    pulses_axis0: int
    pulses_axis2: int
    require_independent_stop: bool
    exercise_dynamic_pause: bool = False


CASES = {
    "independent": AbCase(
        name="独立停止（Q0/Q1先停，Q2/Q3继续）",
        pulses_axis0=100_000,
        pulses_axis2=-200_000,
        require_independent_stop=True,
    ),
    "simultaneous": AbCase(
        name="等长双AB并发完成（相同目标周期数）",
        pulses_axis0=-200_000,
        pulses_axis2=200_000,
        require_independent_stop=False,
    ),
    "dynamic": AbCase(
        name="双AB变频与单组PAUSE/RESUME",
        # At 9600bps the two START transactions, consistent status reads and
        # two atomic frequency writes can consume most of a 300000-cycle
        # 100kHz job. Keep enough motion time for the complete dynamic
        # sequence instead of racing normal completion.
        pulses_axis0=600_000,
        pulses_axis2=-600_000,
        require_independent_stop=False,
        exercise_dynamic_pause=True,
    ),
}


def send_ab_call(
    client: RtuClient, sequence: int, axis: int, operation: int
) -> 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, AB_S2_SET)
    request[12] = axis
    request[13] = AB_OUTPUT_MODE
    request[14] = operation
    client.write_multiple(CALL_REQUEST, request)
    return wait_response(client, CALL_RESPONSE, 12, sequence)


def wait_command_applied(
    client: RtuClient, axis: int, sequence: int, timeout: float = 3.0
) -> dict[str, int]:
    deadline = time.monotonic() + timeout
    latest: dict[str, int] | None = None
    while time.monotonic() < deadline:
        latest = read_axis_status(client, axis)
        if latest["last_sequence"] == sequence:
            if latest["last_result"] != RESULT_OK:
                raise RuntimeError(
                    f"轴{axis}命令#{sequence}执行失败：{latest}"
                )
            return latest
    raise RuntimeError(f"等待轴{axis}命令#{sequence}执行超时：{latest}")


def wait_axis_state(
    client: RtuClient,
    axis: int,
    expected: set[int],
    timeout: float = 4.0,
) -> dict[str, int]:
    deadline = time.monotonic() + timeout
    latest: dict[str, int] | None = None
    while time.monotonic() < deadline:
        latest = read_axis_status(client, axis)
        if latest["error"] != 0:
            raise RuntimeError(f"轴{axis}等待状态时进入错误：{latest}")
        if latest["state"] in expected:
            return latest
    raise RuntimeError(
        f"等待轴{axis}状态{sorted(expected)}超时，最后状态={latest}"
    )


def write_live_frequency(client: RtuClient, frequency_hz: int) -> None:
    words = [0, 0]
    put_u32(words, 0, frequency_hz)
    for axis in AB_OWNERS:
        # One FC16 writes the complete signed INT32 target atomically.
        client.write_multiple(S0_BASES[axis] + 10, words)


def wait_live_frequency(
    client: RtuClient, frequency_hz: int, timeout: float = 3.0
) -> dict[int, dict[str, int]]:
    deadline = time.monotonic() + timeout
    latest: dict[int, dict[str, int]] = {}
    while time.monotonic() < deadline:
        latest = {axis: read_axis_status(client, axis) for axis in AB_OWNERS}
        if all(
            status["state"] in RUNNING_STATES
            and status["target_frequency"] == frequency_hz
            and status["current_frequency"] == frequency_hz
            for status in latest.values()
        ):
            return latest
    raise RuntimeError(
        f"双AB未稳定到{frequency_hz}Hz，最后状态={latest}"
    )


def exercise_dynamic_pause_resume(
    client: RtuClient, sequence: int
) -> int:
    write_live_frequency(client, TEST_LOW_FREQUENCY_HZ)
    slowed = wait_live_frequency(client, TEST_LOW_FREQUENCY_HZ)
    time.sleep(0.05)
    slowed_again = {
        axis: read_axis_status(client, axis) for axis in AB_OWNERS
    }
    if any(
        abs(slowed_again[axis]["task_pulses"])
        <= abs(slowed[axis]["task_pulses"])
        for axis in AB_OWNERS
    ):
        raise RuntimeError(f"双AB降频后计数未继续增长：{slowed_again}")

    write_live_frequency(client, TEST_FREQUENCY_HZ)
    wait_live_frequency(client, TEST_FREQUENCY_HZ)

    response = send_command(client, sequence, 0, CMD_PAUSE, 0)
    check_result(response, 4, RESULT_QUEUED, "轴0 AB PAUSE")
    wait_command_applied(client, 0, sequence)
    sequence += 1
    paused = wait_axis_state(client, 0, {STATE_PAUSED})
    other_before = read_axis_status(client, 2)
    time.sleep(0.05)
    paused_again = read_axis_status(client, 0)
    other_after = read_axis_status(client, 2)
    if (
        paused_again["state"] != STATE_PAUSED
        or paused_again["physical_pulses"] != paused["physical_pulses"]
        or abs(other_after["task_pulses"])
        <= abs(other_before["task_pulses"])
    ):
        raise RuntimeError(
            "AB PAUSE独立性失败："
            f"paused={paused}, paused_again={paused_again}, "
            f"other_before={other_before}, other_after={other_after}"
        )

    response = send_command(client, sequence, 0, CMD_RESUME, 0)
    check_result(response, 4, RESULT_QUEUED, "轴0 AB RESUME")
    wait_command_applied(client, 0, sequence)
    sequence += 1
    resumed = wait_axis_state(client, 0, RUNNING_STATES)
    resume_deadline = time.monotonic() + 1.0
    while time.monotonic() < resume_deadline:
        resumed_again = read_axis_status(client, 0)
        if abs(resumed_again["task_pulses"]) > abs(resumed["task_pulses"]):
            break
    else:
        raise RuntimeError(f"AB RESUME后计数未恢复：{resumed_again}")
    print("动态控制已通过：双AB 100k→50k→100k，Q0/Q1在00边界暂停并恢复")
    return sequence


def prepare_job(client: RtuClient, axis: int, signed_pulses: int) -> None:
    s0 = [0] * 20
    put_u32(s0, 0, 1)
    put_u32(s0, 10, TEST_FREQUENCY_HZ)
    put_u32(s0, 12, signed_pulses)
    client.write_multiple(S0_BASES[axis], s0)
    client.write_multiple(S1_BASES[axis], [0] * 4)


def validate_final_status(
    axis: int,
    status: dict[str, int],
    signed_pulses: int,
    physical_baseline: int,
) -> None:
    expected_physical = abs(signed_pulses)
    checks = {
        "状态": (status["state"], STATE_COMPLETED),
        "错误": (status["error"], 0),
        "执行结果": (status["last_result"], RESULT_OK),
        "逻辑位置": (status["logical_position"], signed_pulses),
        "任务周期数": (status["task_pulses"], signed_pulses),
        "物理周期增量": (
            status["physical_pulses"] - physical_baseline,
            expected_physical,
        ),
    }
    bad = {name: value for name, value in checks.items() if value[0] != value[1]}
    if bad:
        raise RuntimeError(f"轴{axis} AB最终状态不正确：{bad}；状态={status}")


def run_case(
    client: RtuClient, case: AbCase, sequence: int
) -> tuple[int, int]:
    print(f"\n开始用例：{case.name}")
    signed_targets = {0: case.pulses_axis0, 2: case.pulses_axis2}

    for axis in AB_OWNERS:
        response = send_command(client, sequence, axis, CMD_SET_POSITION, 0)
        check_result(response, 4, RESULT_QUEUED, f"轴{axis} SET_POSITION")
        wait_command_applied(client, axis, sequence)
        sequence += 1
        prepare_job(client, axis, signed_targets[axis])

    for axis in AB_OWNERS:
        response = send_ab_call(client, sequence, axis, CALL_COMMIT)
        check_result(response, 3, RESULT_OK, f"轴{axis} AB COMMIT")
        if response[11] != 1:
            raise RuntimeError(f"轴{axis} AB COMMIT未建立有效快照")
        sequence += 1

    baselines = {
        axis: read_axis_status(client, axis)["physical_pulses"]
        for axis in AB_OWNERS
    }
    started = time.monotonic()
    for axis in AB_OWNERS:
        response = send_ab_call(client, sequence, axis, CALL_START)
        check_result(response, 3, RESULT_QUEUED, f"轴{axis} AB START")
        sequence += 1

    running: dict[int, dict[str, int]] = {}
    deadline = time.monotonic() + 4.0
    while time.monotonic() < deadline:
        running = {axis: read_axis_status(client, axis) for axis in AB_OWNERS}
        if all(item["state"] in RUNNING_STATES for item in running.values()):
            break
    else:
        raise RuntimeError(f"双AB未同时进入运行态：{running}")
    for axis, status in running.items():
        if status["counter_mode"] != 1:
            raise RuntimeError(f"轴{axis}未取得AB硬件计数器：{status}")
        if status["current_frequency"] != TEST_FREQUENCY_HZ:
            raise RuntimeError(f"轴{axis}未达到100kHz：{status}")
    print("运行期租约正确：Q0/Q1→TIM9，Q2/Q3→TIM12，双组均为硬件计数")

    if case.exercise_dynamic_pause:
        sequence = exercise_dynamic_pause_resume(client, sequence)

    deadline = time.monotonic() + 8.0
    latest = running
    independent_stop_seen = False
    while time.monotonic() < deadline:
        latest = {axis: read_axis_status(client, axis) for axis in AB_OWNERS}
        if (
            case.require_independent_stop
            and not independent_stop_seen
            and latest[0]["state"] == STATE_COMPLETED
            and latest[2]["state"] in RUNNING_STATES
        ):
            q2_before = abs(latest[2]["task_pulses"])
            q0_physical = latest[0]["physical_pulses"]
            time.sleep(0.05)
            stopped_again = read_axis_status(client, 0)
            running_again = read_axis_status(client, 2)
            independent_stop_seen = (
                stopped_again["state"] == STATE_COMPLETED
                and stopped_again["physical_pulses"] == q0_physical
                and running_again["state"] in RUNNING_STATES
                and abs(running_again["task_pulses"]) > q2_before
            )
            if independent_stop_seen:
                print("独立停止已观测：Q0/Q1保持低且计数冻结，Q2/Q3继续计数")
        if all(item["state"] == STATE_COMPLETED for item in latest.values()):
            break
    else:
        raise RuntimeError(f"等待双AB完成超时：{latest}")
    if case.require_independent_stop and not independent_stop_seen:
        raise RuntimeError("未观测到第一组停止后第二组继续运行的独立停止窗口")

    for axis in AB_OWNERS:
        validate_final_status(
            axis, latest[axis], signed_targets[axis], baselines[axis]
        )
    elapsed = time.monotonic() - started
    print(
        f"用例 PASS：Q0/Q1={abs(case.pulses_axis0)}完整AB周期，"
        f"Q2/Q3={abs(case.pulses_axis2)}完整AB周期，耗时{elapsed:.3f}s"
    )
    return sequence, int(elapsed * 1_000)


def validate_dwt(client: RtuClient, header: list[int]) -> None:
    core_clock_hz = get_u32(header, 5)
    if core_clock_hz == 0 or header[7] != PERFORMANCE_VERSION:
        raise RuntimeError(
            f"P16诊断头无效：clock={core_clock_hz}, version={header[7]}"
        )
    performance = client.read_holding(PERFORMANCE_BASE, 8)
    performance_again = client.read_holding(PERFORMANCE_BASE, 8)
    if performance != performance_again:
        performance = performance_again
    ab_gate_cycles = get_u32(
        client.read_holding(AB_GATE_PERFORMANCE_BASE, 2), 0
    )
    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],
        "AB末周期快速门控": ab_gate_cycles,
    }
    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,
        # Gate must finish within one 100kHz quarter-period (2.5us).
        "AB末周期快速门控": core_clock_hz // 400_000,
    }
    print("\nP16/AB DWT最坏执行时间：")
    for name, cycles in values.items():
        budget = budgets.get(name)
        budget_text = f"预算<{budget}" if budget is not None else "观测项"
        print(
            f"  {name:<18} {cycles:8d} cycles  "
            f"{cycles * 1_000_000.0 / core_clock_hz:8.3f}us  {budget_text}"
        )
    missing = [name for name, cycles in values.items() if cycles == 0]
    if missing:
        raise RuntimeError("DWT路径未实际执行：" + "、".join(missing))
    overruns = {
        name: (values[name], budget)
        for name, budget in budgets.items()
        if values[name] >= budget
    }
    if overruns:
        raise RuntimeError(f"AB实时预算超限：{overruns}")


def print_logic_analyzer_acceptance(selected: list[str]) -> None:
    print("\n逻辑分析仪验收（CH0~CH3=Q0~Q3，建议100MS/s，四通道同步）：")
    if selected == ["independent"]:
        print("  上升沿：CH0=CH1=100000，CH2=CH3=200000。")
    elif selected == ["simultaneous"]:
        print("  上升沿：CH0=CH1=CH2=CH3=200000。")
    elif selected == ["dynamic"]:
        print("  上升沿：CH0=CH1=CH2=CH3=600000（含Q0/Q1暂停窗口）。")
    else:
        print(
            "  连续采全部用例时累计上升沿：CH0=CH1=900000，"
            "CH2=CH3=1000000；精确逐用例验收建议分别用 --case 采集。"
        )
    print(
        "  100kHz区间周期10.000us/高宽5.000us；dynamic的50kHz区间"
        "周期20.000us/高宽10.000us；无<1us窄脉冲。"
    )
    print("  正向：00→10→11→01→00；反向：00→01→11→10→00。")
    print(
        "  独立停止用例：Q0/Q1正向且先停，Q2/Q3反向并继续；"
        "等长用例方向相反。"
    )
    print(
        "  dynamic用例：调频前后保持严格±90°且无额外边沿；Q0/Q1只在00"
        "边界进入低电平暂停，Q2/Q3连续运行，RESUME从00重建相序。"
    )
    print(
        "  每组首跳前必须为00；A/B首个上升沿相隔2.50us（建议容差±0.05us），"
        "不得近似同时上升。"
    )
    print("  末周期必须完整回到00，停止后至少1ms全低、无残余边沿。")


def main() -> int:
    parser = argparse.ArgumentParser(
        description="PLSR 双AB 100kHz硬件计数、独立停止与快速门控压力测试"
    )
    parser.add_argument("--port", help="串口，例如COM5；只有一个串口时可省略")
    parser.add_argument("--baud", type=int, default=9600)
    parser.add_argument("--slave", type=int, default=1)
    parser.add_argument(
        "--case",
        choices=("all", "independent", "simultaneous", "dynamic"),
        default="all",
        help="默认依次执行独立停止、等长并发、动态调频/暂停三个用例",
    )
    args = parser.parse_args()

    selected = (
        ["independent", "simultaneous", "dynamic"]
        if args.case == "all"
        else [args.case]
    )
    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"AB控制窗口未就绪：{header}；请烧录当前固件并硬复位"
            )
        if header[7] != PERFORMANCE_VERSION:
            raise RuntimeError(
                f"AB测试要求性能统计V{PERFORMANCE_VERSION}，当前V{header[7]}"
            )
        if AXIS_STATUS_WORDS != 48:
            raise RuntimeError("上位机轴状态结构版本不匹配")
        print("双AB测试就绪：K3，100kHz，Q0/Q1与Q2/Q3")

        sequence = max(1, (time.monotonic_ns() >> 20) & 0x7FFFFFFF)
        for index, name in enumerate(selected):
            sequence, _elapsed_ms = run_case(client, CASES[name], sequence)
            if index + 1 < len(selected):
                time.sleep(0.25)

        validate_dwt(client, header)
        print_logic_analyzer_acceptance(selected)
        print(
            "\n内部状态与实时预算自动检查PASS；这不代表端子波形通过，"
            "最终结论必须以逻辑分析仪四通道验收为准。"
        )
    return 0


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