"""esp32_protocol.py — Binary frame codec for Jetson↔ESP32-S3 communication.

Issue #119: defines the binary serial protocol between the Jetson Orin Nano Super and the
ESP32-S3 ESP32-S3 BALANCE over USB CDC @ 921600 baud.

Frame layout (all multi-byte fields are big-endian):
  ┌──────┬──────┬──────┬──────────────────┬───────────┬──────┐
  │  STX │ TYPE │  LEN │     PAYLOAD      │  CRC16    │  ETX │
  │ 0x02 │  1B  │  1B  │    LEN bytes     │   2B BE   │ 0x03 │
  └──────┴──────┴──────┴──────────────────┴───────────┴──────┘

CRC16 covers: TYPE + LEN + PAYLOAD (not STX, ETX, or CRC bytes themselves).
CRC algorithm: CCITT-16, polynomial=0x1021, init=0xFFFF, no final XOR.

Command types (Jetson → ESP32-S3):
  0x01  HEARTBEAT   — no payload              (len=0)
  0x02  SPEED_STEER — int16 speed + int16 steer (len=4)  range: -1000..+1000
  0x03  ARM         — uint8 (0=disarm, 1=arm)  (len=1)
  0x04  SET_MODE    — uint8 mode               (len=1)
  0x05  PID_UPDATE  — float32 kp + ki + kd     (len=12)

Telemetry types (ESP32-S3 → Jetson):
  0x10  IMU         — int16×6: pitch,roll,yaw (×100 deg), ax,ay,az (×100 m/s²) (len=12)
  0x11  BATTERY     — uint16 voltage_mv + int16 current_ma + uint8 soc_pct (len=5)
  0x12  MOTOR_RPM   — int16 left_rpm + int16 right_rpm     (len=4)
  0x13  ARM_STATE   — uint8 state + uint8 error_flags       (len=2)
  0x14  ERROR       — uint8 error_code + uint8 subcode      (len=2)

Usage:
  # Encoding (Jetson → ESP32-S3)
  frame = encode_speed_steer(300, -150)
  ser.write(frame)

  # Decoding (ESP32-S3 → Jetson), one byte at a time
  parser = FrameParser()
  for byte in incoming_bytes:
      result = parser.feed(byte)
      if result is not None:
          handle_frame(result)
"""

from __future__ import annotations

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

# ── Frame constants ───────────────────────────────────────────────────────────

STX = 0x02
ETX = 0x03
MAX_PAYLOAD_LEN = 64       # hard limit; any frame larger is corrupt


# ── Command / telemetry type codes ────────────────────────────────────────────

class CmdType(IntEnum):
    HEARTBEAT   = 0x01
    SPEED_STEER = 0x02
    ARM         = 0x03
    SET_MODE    = 0x04
    PID_UPDATE  = 0x05


class TelType(IntEnum):
    IMU       = 0x10
    BATTERY   = 0x11
    MOTOR_RPM = 0x12
    ARM_STATE = 0x13
    ERROR     = 0x14


# ── Parsed telemetry objects ──────────────────────────────────────────────────

@dataclass
class ImuFrame:
    pitch_deg: float      # degrees (positive = forward tilt)
    roll_deg:  float
    yaw_deg:   float
    accel_x:   float      # m/s²
    accel_y:   float
    accel_z:   float


@dataclass
class BatteryFrame:
    voltage_mv:  int      # millivolts (e.g. 11100 = 11.1 V)
    current_ma:  int      # milliamps (negative = charging)
    soc_pct:     int      # state of charge 0–100 (from ESP32-S3 fuel gauge or lookup)


@dataclass
class MotorRpmFrame:
    left_rpm:   int
    right_rpm:  int


@dataclass
class ArmStateFrame:
    state:       int      # 0=DISARMED 1=ARMED 2=TILT_FAULT
    error_flags: int      # bitmask


@dataclass
class ErrorFrame:
    error_code: int
    subcode:    int


# Union type for decoded results
TelemetryFrame = ImuFrame | BatteryFrame | MotorRpmFrame | ArmStateFrame | ErrorFrame


# ── CRC16 CCITT ───────────────────────────────────────────────────────────────

def _crc16_ccitt(data: bytes, init: int = 0xFFFF) -> int:
    """CRC16-CCITT: polynomial 0x1021, init 0xFFFF, no final XOR."""
    crc = init
    for byte in data:
        crc ^= byte << 8
        for _ in range(8):
            if crc & 0x8000:
                crc = (crc << 1) ^ 0x1021
            else:
                crc <<= 1
        crc &= 0xFFFF
    return crc


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

def _build_frame(cmd_type: int, payload: bytes) -> bytes:
    """Assemble a complete binary frame with CRC16."""
    assert len(payload) <= MAX_PAYLOAD_LEN, "Payload too large"
    length = len(payload)
    header = bytes([cmd_type, length])
    crc = _crc16_ccitt(header + payload)
    return bytes([STX, cmd_type, length]) + payload + struct.pack(">H", crc) + bytes([ETX])


def encode_heartbeat() -> bytes:
    """HEARTBEAT frame — no payload."""
    return _build_frame(CmdType.HEARTBEAT, b"")


def encode_speed_steer(speed: int, steer: int) -> bytes:
    """SPEED_STEER frame — int16 speed + int16 steer, both in -1000..+1000."""
    speed = max(-1000, min(1000, int(speed)))
    steer = max(-1000, min(1000, int(steer)))
    return _build_frame(CmdType.SPEED_STEER, struct.pack(">hh", speed, steer))


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


def encode_set_mode(mode: int) -> bytes:
    """SET_MODE frame — mode byte."""
    return _build_frame(CmdType.SET_MODE, struct.pack("B", mode & 0xFF))


def encode_pid_update(kp: float, ki: float, kd: float) -> bytes:
    """PID_UPDATE frame — three float32 values."""
    return _build_frame(CmdType.PID_UPDATE, struct.pack(">fff", kp, ki, kd))


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

class ParserState(IntEnum):
    WAIT_STX  = 0
    WAIT_TYPE = 1
    WAIT_LEN  = 2
    PAYLOAD   = 3
    CRC_HI    = 4
    CRC_LO    = 5
    WAIT_ETX  = 6


class ParseError(Exception):
    pass


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

    Feed individual bytes via feed(); returns a decoded TelemetryFrame (or raw
    bytes tuple) when a complete valid frame is received.

    Thread-safety: single-threaded — wrap in a lock if shared across threads.

    Usage::
        parser = FrameParser()
        for b in incoming:
            result = parser.feed(b)
            if result is not None:
                process(result)
    """

    def __init__(self) -> None:
        self._state    = ParserState.WAIT_STX
        self._type     = 0
        self._length   = 0
        self._payload  = bytearray()
        self._crc_rcvd = 0
        self.frames_ok    = 0
        self.frames_error = 0

    def reset(self) -> None:
        """Reset parser to initial state (call after error or port reconnect)."""
        self._state   = ParserState.WAIT_STX
        self._payload = bytearray()

    def feed(self, byte: int) -> Optional[TelemetryFrame | tuple]:
        """Process one byte. Returns decoded frame on success, None otherwise.

        On CRC error, increments frames_error and resets. The return value on
        success is a dataclass (ImuFrame, BatteryFrame, etc.) or a
        (type_code, raw_payload) tuple for unknown type codes.
        """
        s = self._state

        if s == ParserState.WAIT_STX:
            if byte == STX:
                self._state = ParserState.WAIT_TYPE
            return None

        if s == ParserState.WAIT_TYPE:
            self._type  = byte
            self._state = ParserState.WAIT_LEN
            return None

        if s == ParserState.WAIT_LEN:
            self._length  = byte
            self._payload = bytearray()
            if self._length > MAX_PAYLOAD_LEN:
                # Corrupt frame — too big; reset
                self.frames_error += 1
                self.reset()
                return None
            if self._length == 0:
                self._state = ParserState.CRC_HI
            else:
                self._state = ParserState.PAYLOAD
            return None

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

        if s == ParserState.CRC_HI:
            self._crc_rcvd = byte << 8
            self._state    = ParserState.CRC_LO
            return None

        if s == ParserState.CRC_LO:
            self._crc_rcvd |= byte
            self._state     = ParserState.WAIT_ETX
            return None

        if s == ParserState.WAIT_ETX:
            self.reset()   # always reset so we look for next STX
            if byte != ETX:
                self.frames_error += 1
                return None

            # Verify CRC
            crc_data = bytes([self._type, self._length]) + self._payload
            expected = _crc16_ccitt(crc_data)
            if expected != self._crc_rcvd:
                self.frames_error += 1
                return None

            # Decode
            self.frames_ok += 1
            return _decode_telemetry(self._type, bytes(self._payload))

        # Should never reach here
        self.reset()
        return None


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

def _decode_telemetry(type_code: int, payload: bytes) -> Optional[TelemetryFrame | tuple]:
    """Decode a validated telemetry payload into a typed dataclass."""
    try:
        if type_code == TelType.IMU:
            if len(payload) < 12:
                return None
            p, r, y, ax, ay, az = struct.unpack_from(">hhhhhh", payload)
            return ImuFrame(
                pitch_deg=p  / 100.0,
                roll_deg=r   / 100.0,
                yaw_deg=y    / 100.0,
                accel_x=ax   / 100.0,
                accel_y=ay   / 100.0,
                accel_z=az   / 100.0,
            )

        if type_code == TelType.BATTERY:
            if len(payload) < 5:
                return None
            v_mv, i_ma, soc = struct.unpack_from(">HhB", payload)
            return BatteryFrame(voltage_mv=v_mv, current_ma=i_ma, soc_pct=soc)

        if type_code == TelType.MOTOR_RPM:
            if len(payload) < 4:
                return None
            left, right = struct.unpack_from(">hh", payload)
            return MotorRpmFrame(left_rpm=left, right_rpm=right)

        if type_code == TelType.ARM_STATE:
            if len(payload) < 2:
                return None
            state, flags = struct.unpack_from("BB", payload)
            return ArmStateFrame(state=state, error_flags=flags)

        if type_code == TelType.ERROR:
            if len(payload) < 2:
                return None
            code, sub = struct.unpack_from("BB", payload)
            return ErrorFrame(error_code=code, subcode=sub)

    except struct.error:
        return None

    # Unknown telemetry type — return raw
    return (type_code, payload)
