#!/usr/bin/env python3
"""读取 FANUC J519 允许速度/加速度/jerk 上限表。

本脚本只做 UDP packet type=3 查询，不启动运行时服务，也不写入机器人参数。
响应中的 limit 值保留为控制器原始单位：速度 deg/s、加速度 deg/s^2、jerk deg/s^3。
"""

from __future__ import annotations

import argparse
import csv
import socket
import struct
import sys
from dataclasses import dataclass
from pathlib import Path


LIMIT_TYPES = {
    0: ("velocity", "deg/s"),
    1: ("acceleration", "deg/s^2"),
    2: ("jerk", "deg/s^3"),
}


@dataclass(frozen=True)
class LimitTable:
    """保存单个轴、单种限制类型的 20 档 no-payload / max-payload 上限表。"""

    axis: int
    limit_type: int
    vmax_mm_s: int
    intermediate_check_s: int
    no_payload: tuple[float, ...]
    max_payload: tuple[float, ...]


def pack_request(axis: int, limit_type: int) -> bytes:
    """按手册 Table 5 封装允许上限表查询请求。"""

    return struct.pack(">IIII", 3, 1, axis, limit_type)


def parse_response(packet: bytes, expected_axis: int, expected_type: int) -> LimitTable:
    """按手册 Table 6 解析允许上限表响应。"""

    if len(packet) != 184:
        raise ValueError(f"响应长度应为 184 字节，实际为 {len(packet)} 字节。")

    packet_type, version, axis, limit_type, vmax, check_time = struct.unpack_from(">IIIIII", packet, 0)
    if packet_type != 3 or version != 1:
        raise ValueError(f"响应头不正确：packet_type={packet_type}, version={version}。")
    if axis != expected_axis or limit_type != expected_type:
        raise ValueError(
            f"响应与请求不匹配：期望 axis={expected_axis}, type={expected_type}，"
            f"实际 axis={axis}, type={limit_type}。"
        )

    values = struct.unpack_from(">40f", packet, 24)
    return LimitTable(
        axis=axis,
        limit_type=limit_type,
        vmax_mm_s=vmax,
        intermediate_check_s=check_time,
        no_payload=tuple(values[:20]),
        max_payload=tuple(values[20:]),
    )


def query_table(sock: socket.socket, robot_ip: str, port: int, axis: int, limit_type: int) -> LimitTable:
    """发送单个查询请求并等待对应响应。"""

    sock.sendto(pack_request(axis, limit_type), (robot_ip, port))
    packet, address = sock.recvfrom(2048)
    if address[0] != robot_ip:
        raise ValueError(f"收到非目标机器人响应：{address[0]}:{address[1]}。")
    return parse_response(packet, axis, limit_type)


def iter_requested_tables(
    robot_ip: str,
    port: int,
    timeout_s: float,
    axes: range,
    limit_types: tuple[int, ...],
) -> list[LimitTable]:
    """按轴和限制类型顺序读取所有上限表。"""

    tables: list[LimitTable] = []
    with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as sock:
        sock.settimeout(timeout_s)
        for axis in axes:
            for limit_type in limit_types:
                tables.append(query_table(sock, robot_ip, port, axis, limit_type))
    return tables


def write_csv(path: Path, tables: list[LimitTable]) -> None:
    """把 20 档速度区间和两套 payload 表写成扁平 CSV。"""

    with path.open("w", newline="", encoding="utf-8") as file:
        writer = csv.writer(file)
        writer.writerow(
            [
                "axis",
                "limit_type",
                "unit",
                "vmax_mm_s",
                "intermediate_check_s",
                "speed_bin_index",
                "speed_bin_upper_mm_s",
                "no_payload_limit",
                "max_payload_limit",
            ]
        )
        for table in tables:
            limit_name, unit = LIMIT_TYPES[table.limit_type]
            for index, (no_payload, max_payload) in enumerate(zip(table.no_payload, table.max_payload), start=1):
                writer.writerow(
                    [
                        table.axis,
                        limit_name,
                        unit,
                        table.vmax_mm_s,
                        table.intermediate_check_s,
                        index,
                        table.vmax_mm_s * index / 20.0,
                        f"{no_payload:.9g}",
                        f"{max_payload:.9g}",
                    ]
                )


def print_summary(tables: list[LimitTable]) -> None:
    """打印每张表的首末档，便于现场快速确认返回值是否正常。"""

    for table in tables:
        limit_name, unit = LIMIT_TYPES[table.limit_type]
        print(
            f"axis={table.axis} type={limit_name:<12} unit={unit:<8} "
            f"vmax={table.vmax_mm_s}mm/s check={table.intermediate_check_s}s "
            f"no_payload[1]={table.no_payload[0]:.6g} no_payload[20]={table.no_payload[-1]:.6g} "
            f"max_payload[1]={table.max_payload[0]:.6g} max_payload[20]={table.max_payload[-1]:.6g}"
        )


def parse_args() -> argparse.Namespace:
    """解析命令行参数。"""

    parser = argparse.ArgumentParser(
        description="Read FANUC J519 packet type=3 allowable upper limit tables.",
    )
    parser.add_argument("robot_ip", help="机器人控制柜 IP，例如 192.168.10.11。")
    parser.add_argument("--port", type=int, default=60015, help="J519 UDP 端口，默认 60015。")
    parser.add_argument("--timeout", type=float, default=2.0, help="单次响应超时秒数，默认 2.0。")
    parser.add_argument("--axis-start", type=int, default=1, help="起始轴号，默认 1。")
    parser.add_argument("--axis-end", type=int, default=9, help="结束轴号，默认 9。")
    parser.add_argument(
        "--types",
        default="0,1,2",
        help="限制类型列表：0=velocity, 1=acceleration, 2=jerk；默认 0,1,2。",
    )
    parser.add_argument(
        "--csv",
        type=Path,
        default=Path("analysis/fanuc_allowable_limit_tables.csv"),
        help="CSV 输出路径。",
    )
    return parser.parse_args()


def main() -> int:
    """脚本入口。"""

    args = parse_args()
    limit_types = tuple(int(item.strip()) for item in args.types.split(",") if item.strip())
    unknown_types = [item for item in limit_types if item not in LIMIT_TYPES]
    if unknown_types:
        print(f"未知限制类型：{unknown_types}", file=sys.stderr)
        return 2
    if not 1 <= args.axis_start <= args.axis_end <= 9:
        print("轴范围必须满足 1 <= axis-start <= axis-end <= 9。", file=sys.stderr)
        return 2

    tables = iter_requested_tables(
        args.robot_ip,
        args.port,
        args.timeout,
        range(args.axis_start, args.axis_end + 1),
        limit_types,
    )
    args.csv.parent.mkdir(parents=True, exist_ok=True)
    write_csv(args.csv, tables)
    print_summary(tables)
    print(f"\nCSV 已写入：{args.csv}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())