#!/usr/bin/env python3
"""
balance_protocol.py — CAN message encoding/decoding for the ESP32-S3 BALANCE motor controller
and VESC telemetry.

CAN message layout
------------------
Command frames  (Orin → ESP32-S3 BALANCE / VESC):
  MAMBA_CMD_VELOCITY  0x100  8 bytes  left_speed (f32, m/s) | right_speed (f32, m/s)
  MAMBA_CMD_MODE      0x101  1 byte   mode (0=idle, 1=drive, 2=estop)
  MAMBA_CMD_ESTOP     0x102  1 byte   0x01 = stop

Telemetry frames (ESP32-S3 BALANCE → Orin):
  MAMBA_TELEM_IMU     0x200  24 bytes accel_x, accel_y, accel_z, gyro_x, gyro_y, gyro_z (f32 each)
  MAMBA_TELEM_BATTERY 0x201   8 bytes voltage (f32, V) | current (f32, A)

VESC telemetry frame (VESC → Orin):
  VESC_TELEM_STATE    0x300  16 bytes erpm (f32) | duty (f32) | voltage (f32) | current (f32)

All multi-byte fields are big-endian.

Issue: https://gitea.vayrette.com/seb/saltylab-firmware/issues/674
"""

import struct
from dataclasses import dataclass
from typing import Tuple

# ---------------------------------------------------------------------------
# CAN message IDs
# ---------------------------------------------------------------------------

MAMBA_CMD_VELOCITY: int = 0x100
MAMBA_CMD_MODE: int = 0x101
MAMBA_CMD_ESTOP: int = 0x102

MAMBA_TELEM_IMU: int = 0x200
MAMBA_TELEM_BATTERY: int = 0x201

VESC_TELEM_STATE: int = 0x300
ORIN_CAN_ID_PID_SET: int = 0x305
ORIN_CAN_ID_FC_PID_ACK: int = 0x405

# ---------------------------------------------------------------------------
# Mode constants
# ---------------------------------------------------------------------------

MODE_IDLE: int = 0
MODE_DRIVE: int = 1
MODE_ESTOP: int = 2

# ---------------------------------------------------------------------------
# Data classes for decoded telemetry
# ---------------------------------------------------------------------------


@dataclass
class ImuTelemetry:
    """Decoded IMU telemetry from ESP32-S3 BALANCE (MAMBA_TELEM_IMU)."""

    accel_x: float = 0.0  # m/s²
    accel_y: float = 0.0
    accel_z: float = 0.0
    gyro_x: float = 0.0   # rad/s
    gyro_y: float = 0.0
    gyro_z: float = 0.0


@dataclass
class BatteryTelemetry:
    """Decoded battery telemetry from ESP32-S3 BALANCE (MAMBA_TELEM_BATTERY)."""

    voltage: float = 0.0   # V
    current: float = 0.0   # A


@dataclass
class VescStateTelemetry:
    """Decoded VESC state telemetry (VESC_TELEM_STATE)."""

    erpm: float = 0.0      # electrical RPM
    duty: float = 0.0      # duty cycle [-1.0, 1.0]
    voltage: float = 0.0   # bus voltage, V
    current: float = 0.0   # phase current, A


@dataclass
class PidGains:
    """Balance PID gains (Issue #693)."""
    kp: float = 0.0
    ki: float = 0.0
    kd: float = 0.0


# ---------------------------------------------------------------------------
# Encode helpers
# ---------------------------------------------------------------------------

_FMT_VEL = ">ff"       # 2 × float32, big-endian
_FMT_MODE = ">B"       # 1 × uint8
_FMT_ESTOP = ">B"      # 1 × uint8
_FMT_IMU = ">ffffff"   # 6 × float32
_FMT_BAT = ">ff"       # 2 × float32
_FMT_VESC = ">ffff"    # 4 × float32


def encode_velocity_cmd(left_mps: float, right_mps: float) -> bytes:
    """
    Encode a MAMBA_CMD_VELOCITY payload.

    Parameters
    ----------
    left_mps:  target left wheel speed in m/s  (positive = forward)
    right_mps: target right wheel speed in m/s (positive = forward)

    Returns
    -------
    8-byte big-endian payload suitable for a CAN frame.
    """
    return struct.pack(_FMT_VEL, float(left_mps), float(right_mps))


def encode_mode_cmd(mode: int) -> bytes:
    """
    Encode a MAMBA_CMD_MODE payload.

    Parameters
    ----------
    mode: one of MODE_IDLE (0), MODE_DRIVE (1), MODE_ESTOP (2)

    Returns
    -------
    1-byte payload.
    """
    if mode not in (MODE_IDLE, MODE_DRIVE, MODE_ESTOP):
        raise ValueError(f"Invalid mode {mode!r}; expected 0, 1, or 2")
    return struct.pack(_FMT_MODE, mode)


def encode_estop_cmd(stop: bool = True) -> bytes:
    """
    Encode a MAMBA_CMD_ESTOP payload.

    Parameters
    ----------
    stop: True to assert e-stop, False to clear.

    Returns
    -------
    1-byte payload (0x01 = stop, 0x00 = clear).
    """
    return struct.pack(_FMT_ESTOP, 0x01 if stop else 0x00)


def encode_pid_set_cmd(kp: float, ki: float, kd: float) -> bytes:
    """Encode ORIN_CAN_ID_PID_SET (6 bytes, uint16 BE x3). Issue #693."""
    if kp < 0.0 or ki < 0.0 or kd < 0.0:
        raise ValueError("PID gains must be non-negative")
    return struct.pack(_FMT_PID, round(min(kp,_PID_KP_MAX)*100), round(min(ki,_PID_KI_MAX)*100), round(min(kd,_PID_KD_MAX)*100))


# ---------------------------------------------------------------------------
# Decode helpers
# ---------------------------------------------------------------------------

def decode_imu_telem(data: bytes) -> ImuTelemetry:
    """
    Decode a MAMBA_TELEM_IMU payload.

    Parameters
    ----------
    data: exactly 24 bytes (6 × float32, big-endian).

    Returns
    -------
    ImuTelemetry dataclass instance.

    Raises
    ------
    struct.error if data is the wrong length.
    """
    ax, ay, az, gx, gy, gz = struct.unpack(_FMT_IMU, data)
    return ImuTelemetry(
        accel_x=ax, accel_y=ay, accel_z=az,
        gyro_x=gx, gyro_y=gy, gyro_z=gz,
    )


def decode_battery_telem(data: bytes) -> BatteryTelemetry:
    """
    Decode a MAMBA_TELEM_BATTERY payload.

    Parameters
    ----------
    data: exactly 8 bytes (2 × float32, big-endian).

    Returns
    -------
    BatteryTelemetry dataclass instance.
    """
    voltage, current = struct.unpack(_FMT_BAT, data)
    return BatteryTelemetry(voltage=voltage, current=current)


def decode_vesc_state(data: bytes) -> VescStateTelemetry:
    """
    Decode a VESC_TELEM_STATE payload.

    Parameters
    ----------
    data: exactly 16 bytes (4 × float32, big-endian).

    Returns
    -------
    VescStateTelemetry dataclass instance.
    """
    erpm, duty, voltage, current = struct.unpack(_FMT_VESC, data)
    return VescStateTelemetry(erpm=erpm, duty=duty, voltage=voltage, current=current)


def decode_pid_ack(data: bytes) -> PidGains:
    """Decode ORIN_CAN_ID_FC_PID_ACK (6 bytes). Issue #693."""
    kp_x100, ki_x100, kd_x100 = struct.unpack(_FMT_PID, data)
    return PidGains(kp=kp_x100/100.0, ki=ki_x100/100.0, kd=kd_x100/100.0)
