from __future__ import annotations

import binascii
import socket
import struct
import time
from dataclasses import dataclass
from typing import Iterable

import numpy as np


JOINT_COUNT = 8
PROTOCOL_VERSION = 1

COMMAND_MAGIC = int.from_bytes(b'GUGC', byteorder='little')
TELEMETRY_MAGIC = int.from_bytes(b'GUGT', byteorder='little')

COMMAND_TYPE_SETPOINTS = 1

COMMAND_FLAG_ARM = 1 << 0
COMMAND_FLAG_DISARM = 1 << 1
COMMAND_FLAG_HOLD = 1 << 2
COMMAND_FLAG_ESTOP = 1 << 3
COMMAND_FLAG_CLEAR_FAULTS = 1 << 4

ROBOT_STATE_DISARMED = 0
ROBOT_STATE_ARMED = 1
ROBOT_STATE_ESTOP = 2
ROBOT_STATE_TIMEOUT = 3
ROBOT_STATE_NAMES = {
    ROBOT_STATE_DISARMED: 'DISARMED',
    ROBOT_STATE_ARMED: 'ARMED',
    ROBOT_STATE_ESTOP: 'ESTOP',
    ROBOT_STATE_TIMEOUT: 'TIMEOUT',
}

TELEMETRY_FAULT_MOTOR_MASK = 0x000000FF
TELEMETRY_FAULT_FEEDBACK_TIMEOUT_MASK = 0x0000FF00
TELEMETRY_FAULT_IMU_OFFLINE = 1 << 16
TELEMETRY_FAULT_MAG_OFFLINE = 1 << 17

COMMAND_BODY_FORMAT = '<IBBHIII'
COMMAND_BODY_SIZE = struct.calcsize(COMMAND_BODY_FORMAT) + 5 * JOINT_COUNT * 4
COMMAND_PACKET_SIZE = COMMAND_BODY_SIZE + 4

TELEMETRY_BODY_FORMAT = '<IBBHIIII'
TELEMETRY_BODY_SIZE = struct.calcsize(TELEMETRY_BODY_FORMAT) + 44 * 4
TELEMETRY_PACKET_SIZE = TELEMETRY_BODY_SIZE + 4


def _as_float_array(values: Iterable[float], count: int, *, name: str) -> np.ndarray:
    array = np.asarray(list(values), dtype=np.float32)
    if array.shape != (count,):
        raise ValueError(f'{name} 需要 {count} 个 float，实际得到 {array.shape}')
    return array


def _crc32(payload: bytes) -> int:
    return binascii.crc32(payload) & 0xFFFFFFFF


def _milliseconds_now() -> int:
    return int(time.monotonic() * 1000.0) & 0xFFFFFFFF


@dataclass(slots=True)
class CommandPacket:
    sequence: int
    positions_rad: np.ndarray
    velocities_rad_s: np.ndarray
    kp: np.ndarray
    kd: np.ndarray
    torque_ff_nm: np.ndarray
    flags: int = 0
    monotonic_ms: int | None = None

    def pack(self) -> bytes:
        positions = _as_float_array(self.positions_rad, JOINT_COUNT, name='positions_rad')
        velocities = _as_float_array(self.velocities_rad_s, JOINT_COUNT, name='velocities_rad_s')
        kp = _as_float_array(self.kp, JOINT_COUNT, name='kp')
        kd = _as_float_array(self.kd, JOINT_COUNT, name='kd')
        torque = _as_float_array(self.torque_ff_nm, JOINT_COUNT, name='torque_ff_nm')
        body = struct.pack(
            COMMAND_BODY_FORMAT,
            COMMAND_MAGIC,
            PROTOCOL_VERSION,
            COMMAND_TYPE_SETPOINTS,
            JOINT_COUNT,
            int(self.sequence) & 0xFFFFFFFF,
            _milliseconds_now() if self.monotonic_ms is None else int(self.monotonic_ms) & 0xFFFFFFFF,
            int(self.flags) & 0xFFFFFFFF,
        )
        body += struct.pack(
            '<' + 'f' * 5 * JOINT_COUNT,
            *positions.tolist(),
            *velocities.tolist(),
            *kp.tolist(),
            *kd.tolist(),
            *torque.tolist(),
        )
        return body + struct.pack('<I', _crc32(body))


@dataclass(slots=True)
class TelemetryPacket:
    robot_state: int
    sequence: int
    uptime_ms: int
    last_command_age_ms: int
    fault_mask: int
    joint_position_rad: np.ndarray
    joint_velocity_rad_s: np.ndarray
    joint_torque_nm: np.ndarray
    motor_temperature_c: np.ndarray
    accel_m_s2: np.ndarray
    gyro_rad_s: np.ndarray
    mag_u_t: np.ndarray
    rpy_rad: np.ndarray

    @property
    def roll(self) -> float:
        return float(self.rpy_rad[0])

    @property
    def pitch(self) -> float:
        return float(self.rpy_rad[1])

    @property
    def yaw(self) -> float:
        return float(self.rpy_rad[2])

    @property
    def robot_state_name(self) -> str:
        return ROBOT_STATE_NAMES.get(self.robot_state, f'UNKNOWN({self.robot_state})')

    @property
    def motor_fault_mask(self) -> int:
        return self.fault_mask & TELEMETRY_FAULT_MOTOR_MASK

    @property
    def feedback_timeout_mask(self) -> int:
        return (self.fault_mask & TELEMETRY_FAULT_FEEDBACK_TIMEOUT_MASK) >> 8

    @property
    def imu_offline(self) -> bool:
        return (self.fault_mask & TELEMETRY_FAULT_IMU_OFFLINE) != 0

    @property
    def mag_offline(self) -> bool:
        return (self.fault_mask & TELEMETRY_FAULT_MAG_OFFLINE) != 0

    @property
    def max_motor_temperature_c(self) -> float:
        if self.motor_temperature_c.size == 0:
            return 0.0
        return float(np.max(self.motor_temperature_c))


def unpack_telemetry(packet: bytes) -> TelemetryPacket:
    if len(packet) != TELEMETRY_PACKET_SIZE:
        raise ValueError(f'遥测包长度错误: {len(packet)} != {TELEMETRY_PACKET_SIZE}')
    body = packet[:-4]
    received_crc, = struct.unpack('<I', packet[-4:])
    if _crc32(body) != received_crc:
        raise ValueError('遥测包 CRC 校验失败')

    header_size = struct.calcsize(TELEMETRY_BODY_FORMAT)
    (
        magic,
        version,
        robot_state,
        joint_count,
        sequence,
        uptime_ms,
        last_command_age_ms,
        fault_mask,
    ) = struct.unpack(TELEMETRY_BODY_FORMAT, body[:header_size])
    if magic != TELEMETRY_MAGIC:
        raise ValueError('遥测包 magic 不匹配')
    if version != PROTOCOL_VERSION:
        raise ValueError(f'遥测协议版本不匹配: {version}')
    if joint_count != JOINT_COUNT:
        raise ValueError(f'遥测关节数量不匹配: {joint_count}')

    floats = np.array(
        struct.unpack('<' + 'f' * 44, body[header_size:]),
        dtype=np.float32,
    )
    offset = 0

    def take(count: int) -> np.ndarray:
        nonlocal offset
        result = floats[offset:offset + count].copy()
        offset += count
        return result

    return TelemetryPacket(
        robot_state=robot_state,
        sequence=sequence,
        uptime_ms=uptime_ms,
        last_command_age_ms=last_command_age_ms,
        fault_mask=fault_mask,
        joint_position_rad=take(JOINT_COUNT),
        joint_velocity_rad_s=take(JOINT_COUNT),
        joint_torque_nm=take(JOINT_COUNT),
        motor_temperature_c=take(JOINT_COUNT),
        accel_m_s2=take(3),
        gyro_rad_s=take(3),
        mag_u_t=take(3),
        rpy_rad=take(3),
    )


class UdpRobotLink:
    """主机和 ESP32S3 固件之间的 UDP 通道。"""

    def __init__(
        self,
        *,
        command_host: str,
        command_port: int,
        telemetry_bind_host: str,
        telemetry_port: int,
        timeout_s: float,
    ) -> None:
        self._command_address = (command_host, int(command_port))
        self._socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self._socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self._socket.bind((telemetry_bind_host, int(telemetry_port)))
        self._socket.settimeout(float(timeout_s))

    def send_command(self, command: CommandPacket) -> None:
        self._socket.sendto(command.pack(), self._command_address)

    def receive_telemetry(self) -> TelemetryPacket:
        while True:
            data, _ = self._socket.recvfrom(2048)
            try:
                return unpack_telemetry(data)
            except ValueError:
                # UDP 可能收到旧包、截断包或非本协议包，部署循环里直接跳过。
                continue

    def close(self) -> None:
        self._socket.close()
