#!/usr/bin/env python3
"""
protocol_defs.py — CAN message ID constants and frame builders/parsers for the
Orin↔ESP32-S3 BALANCE↔VESC integration test suite.

All IDs and payload formats are derived from:
  include/orin_can.h   — Orin↔FC (ESP32-S3 BALANCE) protocol
  include/vesc_can.h   — VESC CAN protocol
  saltybot_can_bridge/balance_protocol.py — existing bridge constants

CAN IDs used in tests
---------------------
Orin → FC (ESP32-S3 BALANCE) commands (standard 11-bit, matching orin_can.h):
  ORIN_CMD_HEARTBEAT  0x300
  ORIN_CMD_DRIVE      0x301    int16 speed (−1000..+1000), int16 steer (−1000..+1000)
  ORIN_CMD_MODE       0x302    uint8 mode byte
  ORIN_CMD_ESTOP      0x303    uint8 action (1=ESTOP, 0=CLEAR)

FC (ESP32-S3 BALANCE) → Orin telemetry (standard 11-bit, matching orin_can.h):
  FC_STATUS           0x400    8 bytes (see orin_can_fc_status_t)
  FC_VESC             0x401    8 bytes (see orin_can_fc_vesc_t)
  FC_IMU              0x402    8 bytes
  FC_BARO             0x403    8 bytes

Mamba ↔ VESC internal commands (matching balance_protocol.py):
=======
ESP32-S3 BALANCE ↔ VESC internal commands (matching balance_protocol.py):
>>>>>>> 9aed963 (fix: scrub remaining Mamba references in can_bridge and e2e test protocol files)
  MAMBA_CMD_VELOCITY  0x100    8 bytes  left_mps (f32) | right_mps (f32) big-endian
  MAMBA_CMD_MODE      0x101    1 byte   mode (0=idle,1=drive,2=estop)
  MAMBA_CMD_ESTOP     0x102    1 byte   0x01=stop

VESC STATUS (extended 29-bit, matching vesc_can.h):
  arb_id = (VESC_PKT_STATUS << 8) | vesc_node_id  = (9 << 8) | node_id
  Payload: int32 RPM (BE), int16 current×10 (BE), int16 duty×1000 (BE)
"""

import struct
from typing import Tuple

# ---------------------------------------------------------------------------
# Orin → FC (ESP32-S3 BALANCE) command IDs  (from orin_can.h)
# ---------------------------------------------------------------------------

ORIN_CMD_HEARTBEAT: int = 0x300
ORIN_CMD_DRIVE: int     = 0x301
ORIN_CMD_MODE: int      = 0x302
ORIN_CMD_ESTOP: int     = 0x303

# ---------------------------------------------------------------------------
# FC (ESP32-S3 BALANCE) → Orin telemetry IDs  (from orin_can.h)
# ---------------------------------------------------------------------------

FC_STATUS: int = 0x400
FC_VESC: int   = 0x401
FC_IMU: int    = 0x402
FC_BARO: int   = 0x403

# ---------------------------------------------------------------------------
# Mamba → VESC internal command IDs  (from balance_protocol.py)
=======
# ESP32-S3 BALANCE → VESC internal command IDs  (from balance_protocol.py)
>>>>>>> 9aed963 (fix: scrub remaining Mamba references in can_bridge and e2e test protocol files)
# ---------------------------------------------------------------------------

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

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

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

# ---------------------------------------------------------------------------
# VESC protocol constants  (from vesc_can.h)
# ---------------------------------------------------------------------------

VESC_PKT_STATUS: int    = 9   # STATUS packet type (upper byte of arb_id)
VESC_PKT_SET_RPM: int   = 3   # SET_RPM packet type

VESC_CAN_ID_LEFT: int  = 56
VESC_CAN_ID_RIGHT: int = 68


def VESC_STATUS_ID(vesc_node_id: int) -> int:
    """
    Return the 29-bit extended arbitration ID for a VESC STATUS frame.

    Formula (from vesc_can.h):  arb_id = (VESC_PKT_STATUS << 8) | vesc_node_id
                                        = (9 << 8) | vesc_node_id
    """
    return (VESC_PKT_STATUS << 8) | vesc_node_id


def VESC_SET_RPM_ID(vesc_node_id: int) -> int:
    """
    Return the 29-bit extended arbitration ID for a VESC SET_RPM command.

    Formula:  arb_id = (VESC_PKT_SET_RPM << 8) | vesc_node_id
                     = (3 << 8) | vesc_node_id
    """
    return (VESC_PKT_SET_RPM << 8) | vesc_node_id


# ---------------------------------------------------------------------------
# Frame builders — Orin → FC
# ---------------------------------------------------------------------------

def build_heartbeat(seq: int = 0) -> bytes:
    """Build a HEARTBEAT payload: uint32 sequence counter (big-endian, 4 bytes)."""
    return struct.pack(">I", seq & 0xFFFFFFFF)


def build_drive_cmd(speed: int, steer: int) -> bytes:
    """
    Build an ORIN_CMD_DRIVE payload.

    Parameters
    ----------
    speed: int, −1000..+1000  (mapped directly to int16)
    steer: int, −1000..+1000
    """
    return struct.pack(">hh", int(speed), int(steer))


def build_mode_cmd(mode: int) -> bytes:
    """Build an ORIN_CMD_MODE payload (1 byte)."""
    return struct.pack(">B", mode & 0xFF)


def build_estop_cmd(action: int = 1) -> bytes:
    """Build an ORIN_CMD_ESTOP payload.  action=1 → ESTOP, 0 → CLEAR."""
    return struct.pack(">B", action & 0xFF)


# ---------------------------------------------------------------------------
# Frame builders — Mamba velocity commands (balance_protocol.py encoding)
=======
# Frame builders — ESP32-S3 BALANCE velocity commands (balance_protocol.py encoding)
>>>>>>> 9aed963 (fix: scrub remaining Mamba references in can_bridge and e2e test protocol files)
# ---------------------------------------------------------------------------

def build_velocity_cmd(left_mps: float, right_mps: float) -> bytes:
    """
    Build a MAMBA_CMD_VELOCITY payload (8 bytes, 2 × float32 big-endian).

    Matches encode_velocity_cmd() in balance_protocol.py.
    """
    return struct.pack(">ff", float(left_mps), float(right_mps))


# ---------------------------------------------------------------------------
# Frame builders — FC → Orin telemetry
# ---------------------------------------------------------------------------

def build_fc_status(
    pitch_x10: int = 0,
    motor_cmd: int = 0,
    vbat_mv: int = 24000,
    balance_state: int = 1,
    flags: int = 0,
) -> bytes:
    """
    Build an FC_STATUS (0x400) payload.

    Layout (orin_can_fc_status_t, 8 bytes, big-endian):
      int16  pitch_x10
      int16  motor_cmd
      uint16 vbat_mv
      uint8  balance_state
      uint8  flags  [bit0=estop_active, bit1=armed]
    """
    return struct.pack(
        ">hhHBB",
        int(pitch_x10),
        int(motor_cmd),
        int(vbat_mv) & 0xFFFF,
        int(balance_state) & 0xFF,
        int(flags) & 0xFF,
    )


def build_fc_vesc(
    left_rpm_x10: int = 0,
    right_rpm_x10: int = 0,
    left_current_x10: int = 0,
    right_current_x10: int = 0,
) -> bytes:
    """
    Build an FC_VESC (0x401) payload.

    Layout (orin_can_fc_vesc_t, 8 bytes, big-endian):
      int16 left_rpm_x10
      int16 right_rpm_x10
      int16 left_current_x10
      int16 right_current_x10

    RPM values are RPM / 10 (e.g. 3000 RPM → 300).
    Current values are A × 10 (e.g. 5.5 A → 55).
    """
    return struct.pack(
        ">hhhh",
        int(left_rpm_x10),
        int(right_rpm_x10),
        int(left_current_x10),
        int(right_current_x10),
    )


def build_vesc_status(
    rpm: int = 0,
    current_x10: int = 0,
    duty_x1000: int = 0,
) -> bytes:
    """
    Build a VESC STATUS (packet type 9) payload.

    Layout (from vesc_can.h / VESC FW 6.x, big-endian):
      int32 rpm
      int16 current × 10
      int16 duty × 1000
    Total: 8 bytes.
    """
    return struct.pack(
        ">ihh",
        int(rpm),
        int(current_x10),
        int(duty_x1000),
    )


# ---------------------------------------------------------------------------
# Frame parsers
# ---------------------------------------------------------------------------

def parse_fc_status(data: bytes):
    """
    Parse an FC_STATUS (0x400) payload.

    Returns
    -------
    dict with keys: pitch_x10, motor_cmd, vbat_mv, balance_state, flags,
                    estop_active (bool), armed (bool)
    """
    if len(data) < 8:
        raise ValueError(f"FC_STATUS needs 8 bytes, got {len(data)}")
    pitch_x10, motor_cmd, vbat_mv, balance_state, flags = struct.unpack(
        ">hhHBB", data[:8]
    )
    return {
        "pitch_x10": pitch_x10,
        "motor_cmd": motor_cmd,
        "vbat_mv": vbat_mv,
        "balance_state": balance_state,
        "flags": flags,
        "estop_active": bool(flags & 0x01),
        "armed": bool(flags & 0x02),
    }


def parse_fc_vesc(data: bytes):
    """
    Parse an FC_VESC (0x401) payload.

    Returns
    -------
    dict with keys: left_rpm_x10, right_rpm_x10, left_current_x10,
                    right_current_x10, left_rpm (float), right_rpm (float)
    """
    if len(data) < 8:
        raise ValueError(f"FC_VESC needs 8 bytes, got {len(data)}")
    left_rpm_x10, right_rpm_x10, left_cur_x10, right_cur_x10 = struct.unpack(
        ">hhhh", data[:8]
    )
    return {
        "left_rpm_x10": left_rpm_x10,
        "right_rpm_x10": right_rpm_x10,
        "left_current_x10": left_cur_x10,
        "right_current_x10": right_cur_x10,
        "left_rpm": left_rpm_x10 * 10.0,
        "right_rpm": right_rpm_x10 * 10.0,
    }


def parse_vesc_status(data: bytes):
    """
    Parse a VESC STATUS (packet type 9) payload.

    Returns
    -------
    dict with keys: rpm, current_x10, duty_x1000, current (float), duty (float)
    """
    if len(data) < 8:
        raise ValueError(f"VESC STATUS needs 8 bytes, got {len(data)}")
    rpm, current_x10, duty_x1000 = struct.unpack(">ihh", data[:8])
    return {
        "rpm": rpm,
        "current_x10": current_x10,
        "duty_x1000": duty_x1000,
        "current": current_x10 / 10.0,
        "duty": duty_x1000 / 1000.0,
    }


def parse_velocity_cmd(data: bytes) -> Tuple[float, float]:
    """
    Parse a MAMBA_CMD_VELOCITY payload (8 bytes, 2 × float32 big-endian).

    Returns
    -------
    (left_mps, right_mps)
    """
    if len(data) < 8:
        raise ValueError(f"MAMBA_CMD_VELOCITY needs 8 bytes, got {len(data)}")
    return struct.unpack(">ff", data[:8])
