"""esp32_balance_protocol.py — Binary frame codec for Orin↔ESP32-S3 BALANCE UART/USB comms.

bd-wim1: replaces CANable2/python-can ESP32 interface with USB-CDC serial.
Counterpart firmware bead: bd-66hx (hal's ESP32-S3 BALANCE side).

Frame layout:
  ┌───────┬─────┬──────┬──────────────────┬───────┐
  │  SYNC │ LEN │ TYPE │     PAYLOAD      │  CRC8 │
  │  0xAA │  1B │  1B  │    LEN bytes     │   1B  │
  └───────┴─────┴──────┴──────────────────┴───────┘

  LEN   = number of payload bytes (0–62)
  CRC8  = CRC8-SMBUS (poly=0x07, init=0x00), computed over LEN+TYPE+PAYLOAD

Physical layer
  Device  : /dev/esp32-balance  (CH343 USB-CDC, VID:PID 1a86:55d3)
  Baud    : 460800, 8N1

Command types (Orin → ESP32-S3 BALANCE):
  0x01  CMD_HEARTBEAT  — no payload                               len=0
  0x02  CMD_DRIVE      — int16 speed_units, int16 steer_units     len=4  range −1000..+1000
  0x03  CMD_ESTOP      — uint8 (1=assert, 0=clear)               len=1
  0x04  CMD_ARM        — uint8 (1=arm, 0=disarm)                 len=1
  0x05  CMD_PID        — float32 kp, float32 ki, float32 kd       len=12

Telemetry types (ESP32-S3 BALANCE → Orin):
  0x80  TELEM_STATUS      — int16 pitch_x10, int16 motor_cmd,
                            uint16 vbat_mv, uint8 balance_state, uint8 flags   len=8
  0x81  TELEM_VESC_LEFT   — int32 erpm, uint16 voltage_mv,
                            int16 current_ma, uint16 temp_c_x10                len=10
  0x82  TELEM_VESC_RIGHT  — same layout as TELEM_VESC_LEFT                     len=10
  0xA0  RESP_ACK   — uint8 echoed cmd type                       len=1
  0xA1  RESP_NACK  — uint8 cmd type, uint8 error_code            len=2

Balance state values (TELEM_STATUS.balance_state):
  0  DISARMED
  1  ARMED
  2  TILT_FAULT
  3  ESTOP

NACK error codes:
  0x01  ERR_BAD_CRC
  0x02  ERR_BAD_LEN
  0x03  ERR_ESTOP_ACTIVE
  0x04  ERR_DISARMED

NOTE: Message type IDs coordinated with hal (bd-66hx) on 2026-04-17.
      Confirm against bd-66hx PR before deploying on hardware.
"""

from __future__ import annotations

import struct
from dataclasses import dataclass
from enum import IntEnum
from typing import Optional

# ── Physical layer constants ──────────────────────────────────────────────────

SERIAL_PORT_DEFAULT = "/dev/esp32-balance"
BAUD_RATE           = 460800
MAX_PAYLOAD_LEN     = 62       # hard cap; larger frames are considered corrupt

# ── Framing ───────────────────────────────────────────────────────────────────

SYNC_BYTE = 0xAA


# ── Command type codes (Orin → ESP32-S3 BALANCE) ─────────────────────────────

class CmdType(IntEnum):
    HEARTBEAT = 0x01
    DRIVE     = 0x02
    ESTOP     = 0x03
    ARM       = 0x04
    PID       = 0x05


# ── Telemetry type codes (ESP32-S3 BALANCE → Orin) ───────────────────────────

class TelType(IntEnum):
    STATUS     = 0x80
    VESC_LEFT  = 0x81
    VESC_RIGHT = 0x82
    ACK        = 0xA0
    NACK       = 0xA1


# ── Balance state enum ────────────────────────────────────────────────────────

class BalanceState(IntEnum):
    DISARMED   = 0
    ARMED      = 1
    TILT_FAULT = 2
    ESTOP      = 3


BALANCE_STATE_LABEL = {
    BalanceState.DISARMED:   "DISARMED",
    BalanceState.ARMED:      "ARMED",
    BalanceState.TILT_FAULT: "TILT_FAULT",
    BalanceState.ESTOP:      "ESTOP",
}


# ── Parsed telemetry dataclasses ─────────────────────────────────────────────

@dataclass
class StatusFrame:
    """TELEM_STATUS (0x80): balance controller status at ~10 Hz."""
    pitch_deg:     float   # degrees (positive = forward tilt); raw pitch_x10 / 10.0
    motor_cmd:     int     # ESC command, −1000..+1000
    vbat_mv:       int     # battery voltage, millivolts
    balance_state: int     # BalanceState enum
    flags:         int     # bitmask — bit 0: estop_active, bit 1: heartbeat_timeout


@dataclass
class VescFrame:
    """TELEM_VESC_LEFT (0x81) or TELEM_VESC_RIGHT (0x82): VESC proxied telemetry at ~10 Hz."""
    vesc_id:    int    # 56 = left, 68 = right (maps to CAN IDs from hardware spec)
    erpm:       int    # electrical RPM (signed)
    voltage_mv: int    # bus voltage, millivolts
    current_ma: int    # phase current, milliamps (signed)
    temp_c:     float  # motor temperature, °C; raw temp_c_x10 / 10.0


@dataclass
class AckFrame:
    """RESP_ACK (0xA0): acknowledgement for a successfully parsed command."""
    echoed_type: int


@dataclass
class NackFrame:
    """RESP_NACK (0xA1): rejection; contains original cmd type and error code."""
    cmd_type:   int
    error_code: int


TelemetryFrame = StatusFrame | VescFrame | AckFrame | NackFrame


# ── CRC8-SMBUS ────────────────────────────────────────────────────────────────

def _crc8(data: bytes) -> int:
    """CRC8-SMBUS: polynomial 0x07, init 0x00, no final XOR.

    Coverage: LEN + TYPE + PAYLOAD bytes.
    """
    crc = 0
    for byte in data:
        crc ^= byte
        for _ in range(8):
            if crc & 0x80:
                crc = (crc << 1) ^ 0x07
            else:
                crc <<= 1
        crc &= 0xFF
    return crc


# ── Frame encoder ─────────────────────────────────────────────────────────────

def _build_frame(cmd_type: int, payload: bytes) -> bytes:
    """Build a complete binary frame: [0xAA][LEN][TYPE][PAYLOAD][CRC8]."""
    if len(payload) > MAX_PAYLOAD_LEN:
        raise ValueError(f"Payload length {len(payload)} exceeds max {MAX_PAYLOAD_LEN}")
    length = len(payload)
    crc_data = bytes([length, cmd_type]) + payload
    crc = _crc8(crc_data)
    return bytes([SYNC_BYTE, length, cmd_type]) + payload + bytes([crc])


def encode_heartbeat() -> bytes:
    """CMD_HEARTBEAT — no payload. Send every ≤500 ms to keep watchdog alive."""
    return _build_frame(CmdType.HEARTBEAT, b"")


def encode_drive(speed: int, steer: int) -> bytes:
    """CMD_DRIVE — int16 speed_units + int16 steer_units (big-endian, −1000..+1000).

    Parameters
    ----------
    speed : target longitudinal speed in motor units (positive = forward)
    steer : steering demand in motor units (positive = right turn)
    """
    speed = max(-1000, min(1000, int(speed)))
    steer = max(-1000, min(1000, int(steer)))
    return _build_frame(CmdType.DRIVE, struct.pack(">hh", speed, steer))


def encode_estop(assert_stop: bool) -> bytes:
    """CMD_ESTOP — uint8, 1=assert, 0=clear."""
    return _build_frame(CmdType.ESTOP, struct.pack("B", 1 if assert_stop else 0))


def encode_arm(arm: bool) -> bytes:
    """CMD_ARM — uint8, 1=arm, 0=disarm."""
    return _build_frame(CmdType.ARM, struct.pack("B", 1 if arm else 0))


def encode_pid(kp: float, ki: float, kd: float) -> bytes:
    """CMD_PID — three float32 values, big-endian."""
    if kp < 0.0 or ki < 0.0 or kd < 0.0:
        raise ValueError("PID gains must be non-negative")
    return _build_frame(CmdType.PID, struct.pack(">fff", kp, ki, kd))


# ── Frame decoder (streaming state-machine) ───────────────────────────────────

class _ParseState(IntEnum):
    WAIT_SYNC = 0
    WAIT_LEN  = 1
    WAIT_TYPE = 2
    PAYLOAD   = 3
    WAIT_CRC  = 4


class FrameParser:
    """Byte-by-byte streaming parser for ESP32-S3 BALANCE telemetry frames.

    Feed individual bytes via ``feed()``; returns a decoded dataclass when a
    complete, CRC-valid frame arrives, otherwise returns ``None``.

    Not thread-safe — guard with a lock if accessed from multiple threads.

    Counters
    --------
    frames_ok     : successfully decoded frames
    frames_error  : CRC failures or oversized payloads
    """

    def __init__(self) -> None:
        self.frames_ok    = 0
        self.frames_error = 0
        self._reset()

    def reset(self) -> None:
        """Reset to initial state — call after serial reconnect."""
        self._reset()

    def _reset(self) -> None:
        self._state   = _ParseState.WAIT_SYNC
        self._length  = 0
        self._type    = 0
        self._payload = bytearray()

    def feed(self, byte: int) -> Optional[TelemetryFrame | tuple]:
        """Process one byte.

        Returns a decoded dataclass on success, ``None`` while waiting for more
        bytes, or a ``(type_code, raw_payload)`` tuple for unrecognised types.
        """
        s = self._state

        if s == _ParseState.WAIT_SYNC:
            if byte == SYNC_BYTE:
                self._state = _ParseState.WAIT_LEN
            return None

        if s == _ParseState.WAIT_LEN:
            if byte > MAX_PAYLOAD_LEN:
                self.frames_error += 1
                self._reset()
                return None
            self._length = byte
            self._state  = _ParseState.WAIT_TYPE
            return None

        if s == _ParseState.WAIT_TYPE:
            self._type    = byte
            self._payload = bytearray()
            if self._length == 0:
                self._state = _ParseState.WAIT_CRC
            else:
                self._state = _ParseState.PAYLOAD
            return None

        if s == _ParseState.PAYLOAD:
            self._payload.append(byte)
            if len(self._payload) == self._length:
                self._state = _ParseState.WAIT_CRC
            return None

        if s == _ParseState.WAIT_CRC:
            # Verify CRC before resetting parser state
            crc_data = bytes([self._length, self._type]) + self._payload
            expected = _crc8(crc_data)
            if byte != expected:
                self.frames_error += 1
                self._reset()
                return None
            self.frames_ok += 1
            result = _decode(self._type, bytes(self._payload))
            self._reset()
            return result

        self._reset()
        return None


# ── Telemetry decoder ─────────────────────────────────────────────────────────

# Maps VESC_LEFT/RIGHT type codes to hardware CAN IDs (from memory)
_VESC_ID_MAP = {
    TelType.VESC_LEFT:  56,   # left VESC CAN ID
    TelType.VESC_RIGHT: 68,   # right VESC CAN ID
}


def _decode(type_code: int, payload: bytes) -> Optional[TelemetryFrame | tuple]:
    """Decode a validated telemetry payload into a typed dataclass."""
    try:
        if type_code == TelType.STATUS:
            # int16 pitch_x10, int16 motor_cmd, uint16 vbat_mv, uint8 state, uint8 flags
            if len(payload) < 8:
                return None
            pitch_x10, motor_cmd, vbat_mv, state, flags = struct.unpack_from(">hhHBB", payload)
            return StatusFrame(
                pitch_deg=pitch_x10 / 10.0,
                motor_cmd=motor_cmd,
                vbat_mv=vbat_mv,
                balance_state=state,
                flags=flags,
            )

        if type_code in (TelType.VESC_LEFT, TelType.VESC_RIGHT):
            # int32 erpm, uint16 voltage_mv, int16 current_ma, uint16 temp_c_x10
            if len(payload) < 10:
                return None
            erpm, voltage_mv, current_ma, temp_x10 = struct.unpack_from(">iHhH", payload)
            return VescFrame(
                vesc_id=_VESC_ID_MAP[type_code],
                erpm=erpm,
                voltage_mv=voltage_mv,
                current_ma=current_ma,
                temp_c=temp_x10 / 10.0,
            )

        if type_code == TelType.ACK:
            if len(payload) < 1:
                return None
            return AckFrame(echoed_type=payload[0])

        if type_code == TelType.NACK:
            if len(payload) < 2:
                return None
            return NackFrame(cmd_type=payload[0], error_code=payload[1])

    except struct.error:
        return None

    # Unknown type — return raw for forward-compatibility
    return (type_code, payload)
