"""test_esp32_protocol.py — Unit tests for binary ESP32-S3 frame codec.

Tests:
  - CRC16-CCITT correctness
  - All encoder functions produce correctly structured frames
  - FrameParser decodes valid frames for all telemetry types
  - FrameParser rejects frames with bad CRC
  - FrameParser rejects frames with bad ETX
  - FrameParser resyncs after corrupt data (finds next STX)
  - FrameParser handles streaming byte-by-byte input
  - FrameParser handles oversized payload gracefully
  - Speed/steer clamping in encode_speed_steer
  - Round-trip encode → decode for all known telemetry types

Run with: pytest test/test_esp32_protocol.py -v
"""

from __future__ import annotations

import struct
import pytest
import sys
import os

# ── Path setup (no ROS2 install needed) ──────────────────────────────────────
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

from saltybot_bridge.esp32_protocol import (
    STX, ETX,
    CmdType, TelType,
    ImuFrame, BatteryFrame, MotorRpmFrame, ArmStateFrame, ErrorFrame,
    FrameParser,
    _crc16_ccitt, _build_frame, _decode_telemetry,
    encode_heartbeat, encode_speed_steer, encode_arm,
    encode_set_mode, encode_pid_update,
)


# ── Helpers ───────────────────────────────────────────────────────────────────

def _parse_all(raw: bytes):
    """Feed all bytes into a fresh FrameParser, return list of decoded frames."""
    parser = FrameParser()
    results = []
    for b in raw:
        r = parser.feed(b)
        if r is not None:
            results.append(r)
    return results, parser


def _build_telem_frame(type_code: int, payload: bytes) -> bytes:
    """Build a raw telemetry frame (same framing as command encoder)."""
    return _build_frame(type_code, payload)


# ═══════════════════════════════════════════════════════════════════
# CRC16 tests
# ═══════════════════════════════════════════════════════════════════

class TestCrc16:
    def test_known_vector_empty(self):
        """CRC of empty data must equal 0xFFFF (init value)."""
        assert _crc16_ccitt(b"") == 0xFFFF

    def test_known_vector_123456789(self):
        """Classic CCITT-16 test vector: '123456789' → 0x29B1."""
        assert _crc16_ccitt(b"123456789") == 0x29B1

    def test_crc_changes_on_bit_flip(self):
        data  = b"\x02\x04\x42"
        flipped = b"\x02\x04\x43"
        assert _crc16_ccitt(data) != _crc16_ccitt(flipped)

    def test_crc_all_zeros(self):
        crc = _crc16_ccitt(bytes(16))
        assert isinstance(crc, int)
        assert 0 <= crc <= 0xFFFF

    def test_crc_single_byte(self):
        c1 = _crc16_ccitt(b"\x01")
        c2 = _crc16_ccitt(b"\x02")
        assert c1 != c2


# ═══════════════════════════════════════════════════════════════════
# Encoder tests
# ═══════════════════════════════════════════════════════════════════

class TestEncoders:
    def test_heartbeat_structure(self):
        frame = encode_heartbeat()
        # STX TYPE LEN CRC_HI CRC_LO ETX  (no payload → 6 bytes)
        assert len(frame) == 6
        assert frame[0] == STX
        assert frame[1] == CmdType.HEARTBEAT
        assert frame[2] == 0          # LEN=0
        assert frame[-1] == ETX

    def test_heartbeat_crc_valid(self):
        frame = encode_heartbeat()
        crc_expected = _crc16_ccitt(bytes([CmdType.HEARTBEAT, 0]))
        crc_actual   = struct.unpack(">H", frame[3:5])[0]
        assert crc_actual == crc_expected

    def test_speed_steer_structure(self):
        frame = encode_speed_steer(300, -150)
        assert frame[0] == STX
        assert frame[1] == CmdType.SPEED_STEER
        assert frame[2] == 4           # LEN=4 (two int16)
        assert frame[-1] == ETX
        assert len(frame) == 10        # 3 header + 4 payload + 2 CRC + 1 ETX

    def test_speed_steer_payload_values(self):
        frame = encode_speed_steer(500, -200)
        speed, steer = struct.unpack(">hh", frame[3:7])
        assert speed == 500
        assert steer == -200

    def test_speed_steer_clamped_max(self):
        frame = encode_speed_steer(9999, 9999)
        speed, steer = struct.unpack(">hh", frame[3:7])
        assert speed == 1000
        assert steer == 1000

    def test_speed_steer_clamped_min(self):
        frame = encode_speed_steer(-9999, -9999)
        speed, steer = struct.unpack(">hh", frame[3:7])
        assert speed == -1000
        assert steer == -1000

    def test_speed_steer_zero(self):
        frame = encode_speed_steer(0, 0)
        speed, steer = struct.unpack(">hh", frame[3:7])
        assert speed == 0
        assert steer == 0

    def test_arm_true(self):
        frame = encode_arm(True)
        assert frame[1] == CmdType.ARM
        assert frame[2] == 1           # LEN=1
        assert frame[3] == 1           # payload = 1 (arm)

    def test_arm_false(self):
        frame = encode_arm(False)
        assert frame[3] == 0           # payload = 0 (disarm)

    def test_set_mode(self):
        frame = encode_set_mode(3)
        assert frame[1] == CmdType.SET_MODE
        assert frame[3] == 3

    def test_set_mode_masked(self):
        frame = encode_set_mode(0x1FF)
        assert frame[3] == 0xFF        # masked to uint8

    def test_pid_update_structure(self):
        frame = encode_pid_update(1.5, 0.01, 0.1)
        assert frame[1] == CmdType.PID_UPDATE
        assert frame[2] == 12          # LEN=12 (three float32)
        kp, ki, kd = struct.unpack(">fff", frame[3:15])
        assert kp == pytest.approx(1.5)
        assert ki == pytest.approx(0.01)
        assert kd == pytest.approx(0.1)

    def test_pid_update_zero(self):
        frame = encode_pid_update(0.0, 0.0, 0.0)
        kp, ki, kd = struct.unpack(">fff", frame[3:15])
        assert kp == pytest.approx(0.0)

    def test_all_frames_end_with_etx(self):
        frames = [
            encode_heartbeat(),
            encode_speed_steer(100, -100),
            encode_arm(True),
            encode_set_mode(1),
            encode_pid_update(2.0, 0.1, 0.05),
        ]
        for f in frames:
            assert f[-1] == ETX, f"Frame for type 0x{f[1]:02X} doesn't end with ETX"

    def test_all_frames_start_with_stx(self):
        for f in [encode_heartbeat(), encode_speed_steer(0, 0)]:
            assert f[0] == STX


# ═══════════════════════════════════════════════════════════════════
# FrameParser tests
# ═══════════════════════════════════════════════════════════════════

class TestFrameParser:

    # ── IMU telemetry ─────────────────────────────────────────────────────────

    def test_decode_imu_frame(self):
        payload = struct.pack(">hhhhhh",
            1250, -300, 900,   # pitch×100, roll×100, yaw×100
            981,  0,   -50,   # ax×100, ay×100, az×100
        )
        raw = _build_telem_frame(TelType.IMU, payload)
        results, parser = _parse_all(raw)
        assert len(results) == 1
        f = results[0]
        assert isinstance(f, ImuFrame)
        assert f.pitch_deg == pytest.approx(12.5)
        assert f.roll_deg  == pytest.approx(-3.0)
        assert f.yaw_deg   == pytest.approx(9.0)
        assert f.accel_x   == pytest.approx(9.81)
        assert f.accel_y   == pytest.approx(0.0)
        assert f.accel_z   == pytest.approx(-0.5)

    def test_decode_imu_zero(self):
        payload = struct.pack(">hhhhhh", 0, 0, 0, 0, 0, 0)
        raw = _build_telem_frame(TelType.IMU, payload)
        results, _ = _parse_all(raw)
        f = results[0]
        assert f.pitch_deg == pytest.approx(0.0)
        assert f.accel_z   == pytest.approx(0.0)

    # ── Battery telemetry ─────────────────────────────────────────────────────

    def test_decode_battery_frame(self):
        payload = struct.pack(">HhB", 11100, -500, 72)  # 11.1V, -500mA (charging), 72%
        raw = _build_telem_frame(TelType.BATTERY, payload)
        results, _ = _parse_all(raw)
        assert len(results) == 1
        f = results[0]
        assert isinstance(f, BatteryFrame)
        assert f.voltage_mv == 11100
        assert f.current_ma == -500
        assert f.soc_pct    == 72

    def test_decode_battery_discharging(self):
        payload = struct.pack(">HhB", 10200, 2000, 45)  # 10.2V, 2A draw, 45%
        raw = _build_telem_frame(TelType.BATTERY, payload)
        results, _ = _parse_all(raw)
        f = results[0]
        assert f.current_ma == 2000     # positive = discharging

    # ── Motor RPM telemetry ───────────────────────────────────────────────────

    def test_decode_motor_rpm(self):
        payload = struct.pack(">hh", 1500, -1500)   # left=1500, right=-1500 (spinning in place)
        raw = _build_telem_frame(TelType.MOTOR_RPM, payload)
        results, _ = _parse_all(raw)
        f = results[0]
        assert isinstance(f, MotorRpmFrame)
        assert f.left_rpm  == 1500
        assert f.right_rpm == -1500

    def test_decode_motor_rpm_zero(self):
        payload = struct.pack(">hh", 0, 0)
        raw = _build_telem_frame(TelType.MOTOR_RPM, payload)
        results, _ = _parse_all(raw)
        assert results[0].left_rpm == 0

    # ── Arm state telemetry ───────────────────────────────────────────────────

    def test_decode_arm_state_armed(self):
        payload = struct.pack("BB", 1, 0)   # ARMED, no error flags
        raw = _build_telem_frame(TelType.ARM_STATE, payload)
        results, _ = _parse_all(raw)
        f = results[0]
        assert isinstance(f, ArmStateFrame)
        assert f.state       == 1
        assert f.error_flags == 0

    def test_decode_arm_state_tilt_fault(self):
        payload = struct.pack("BB", 2, 0x04)   # TILT_FAULT, bit 2 set
        raw = _build_telem_frame(TelType.ARM_STATE, payload)
        results, _ = _parse_all(raw)
        f = results[0]
        assert f.state       == 2
        assert f.error_flags == 0x04

    # ── Error telemetry ───────────────────────────────────────────────────────

    def test_decode_error_frame(self):
        payload = struct.pack("BB", 0x0A, 0x01)   # error 0x0A, sub 0x01
        raw = _build_telem_frame(TelType.ERROR, payload)
        results, _ = _parse_all(raw)
        f = results[0]
        assert isinstance(f, ErrorFrame)
        assert f.error_code == 0x0A
        assert f.subcode    == 0x01

    # ── CRC errors ────────────────────────────────────────────────────────────

    def test_bad_crc_rejected(self):
        raw = bytearray(_build_telem_frame(TelType.BATTERY,
                         struct.pack(">HhB", 11000, 0, 80)))
        raw[-3] ^= 0xFF    # flip CRC_HI byte
        results, parser = _parse_all(bytes(raw))
        assert results == []
        assert parser.frames_error == 1
        assert parser.frames_ok    == 0

    def test_bad_crc_both_bytes(self):
        raw = bytearray(_build_telem_frame(TelType.MOTOR_RPM,
                         struct.pack(">hh", 100, 100)))
        raw[-3] ^= 0xAA
        raw[-2] ^= 0x55
        results, parser = _parse_all(bytes(raw))
        assert results == []
        assert parser.frames_error >= 1

    # ── ETX errors ────────────────────────────────────────────────────────────

    def test_bad_etx_rejected(self):
        raw = bytearray(_build_telem_frame(TelType.ARM_STATE,
                         struct.pack("BB", 1, 0)))
        raw[-1] = 0xFF     # corrupt ETX
        results, parser = _parse_all(bytes(raw))
        assert results == []
        assert parser.frames_error == 1

    # ── Resync after corrupt data ─────────────────────────────────────────────

    def test_resync_after_garbage(self):
        """Parser must resync and decode a valid frame after leading garbage."""
        garbage = b"\xDE\xAD\xBE\xEF\xFF\x00\x42"
        valid   = _build_telem_frame(TelType.ARM_STATE, struct.pack("BB", 1, 0))
        results, parser = _parse_all(garbage + valid)
        assert len(results) == 1
        assert isinstance(results[0], ArmStateFrame)

    def test_two_consecutive_frames(self):
        """Parser must handle two back-to-back valid frames."""
        f1 = _build_telem_frame(TelType.ARM_STATE, struct.pack("BB", 1, 0))
        f2 = _build_telem_frame(TelType.MOTOR_RPM, struct.pack(">hh", 500, 500))
        results, parser = _parse_all(f1 + f2)
        assert len(results) == 2
        assert isinstance(results[0], ArmStateFrame)
        assert isinstance(results[1], MotorRpmFrame)
        assert parser.frames_error == 0

    def test_three_frames_with_garbage_between(self):
        """Corrupt bytes between valid frames must not discard subsequent valid frames."""
        f1 = _build_telem_frame(TelType.ARM_STATE, struct.pack("BB", 1, 0))
        f2 = _build_telem_frame(TelType.BATTERY,   struct.pack(">HhB", 11000, 0, 80))
        junk = b"\xDE\xAD"
        results, parser = _parse_all(f1 + junk + f2)
        assert len(results) == 2

    def test_unknown_type_returns_raw_tuple(self):
        """Unknown telemetry type must be returned as (type_code, payload) tuple."""
        raw = _build_telem_frame(0x99, b"\xAB\xCD")
        results, _ = _parse_all(raw)
        assert len(results) == 1
        type_code, payload = results[0]
        assert type_code == 0x99
        assert payload   == b"\xAB\xCD"

    def test_oversized_payload_rejected(self):
        """A frame claiming LEN > MAX_PAYLOAD_LEN must be silently skipped."""
        # Craft a raw frame with LEN=200 (> 64 limit)
        raw = bytes([STX, TelType.IMU, 200])
        results, parser = _parse_all(raw)
        assert results == []
        assert parser.frames_error == 1

    def test_empty_payload_frame(self):
        """A frame with LEN=0 (heartbeat style) must decode correctly."""
        raw = _build_frame(CmdType.HEARTBEAT, b"")
        results, parser = _parse_all(raw)
        assert len(results) == 1
        type_code, payload = results[0]
        assert type_code == CmdType.HEARTBEAT
        assert payload   == b""
        assert parser.frames_error == 0

    # ── Streaming byte-by-byte ────────────────────────────────────────────────

    def test_stream_one_byte_at_a_time(self):
        """Every byte fed individually must produce the same result as bulk."""
        payload = struct.pack(">HhB", 12000, -100, 90)
        raw = _build_telem_frame(TelType.BATTERY, payload)
        parser = FrameParser()
        results = []
        for b in raw:
            r = parser.feed(b)
            if r is not None:
                results.append(r)
        assert len(results) == 1
        f = results[0]
        assert isinstance(f, BatteryFrame)
        assert f.soc_pct == 90

    # ── Parser reset ──────────────────────────────────────────────────────────

    def test_parser_reset_clears_state(self):
        """After reset(), parser must accept a new valid frame."""
        payload = struct.pack("BB", 0, 0)
        raw = _build_telem_frame(TelType.ARM_STATE, payload)
        parser = FrameParser()
        # Feed partial frame
        for b in raw[:3]:
            parser.feed(b)
        # Reset mid-frame
        parser.reset()
        # Feed full frame
        results = []
        for b in raw:
            r = parser.feed(b)
            if r is not None:
                results.append(r)
        assert len(results) == 1

    def test_frames_ok_counter(self):
        f1 = _build_telem_frame(TelType.ARM_STATE, struct.pack("BB", 1, 0))
        f2 = _build_telem_frame(TelType.ARM_STATE, struct.pack("BB", 0, 0))
        _, parser = _parse_all(f1 + f2)
        assert parser.frames_ok == 2

    def test_frames_error_counter(self):
        raw = bytearray(_build_telem_frame(TelType.ARM_STATE, struct.pack("BB", 1, 0)))
        raw[-3] ^= 0xFF   # corrupt CRC
        _, parser = _parse_all(bytes(raw))
        assert parser.frames_error == 1


# ═══════════════════════════════════════════════════════════════════
# Decode edge cases
# ═══════════════════════════════════════════════════════════════════

class TestDecodeEdgeCases:
    def test_imu_short_payload_returns_none(self):
        result = _decode_telemetry(TelType.IMU, b"\x00\x01")  # only 2 bytes, need 12
        assert result is None

    def test_battery_short_payload_returns_none(self):
        result = _decode_telemetry(TelType.BATTERY, b"\x00")
        assert result is None

    def test_motor_rpm_short_payload_returns_none(self):
        result = _decode_telemetry(TelType.MOTOR_RPM, b"\x00\x01")
        assert result is None

    def test_arm_state_short_payload_returns_none(self):
        result = _decode_telemetry(TelType.ARM_STATE, b"\x01")
        assert result is None

    def test_error_short_payload_returns_none(self):
        result = _decode_telemetry(TelType.ERROR, b"\x01")
        assert result is None


# ═══════════════════════════════════════════════════════════════════
# Round-trip tests (encode command → re-parse frame body)
# ═══════════════════════════════════════════════════════════════════

class TestRoundTrip:
    """Verify encoder output satisfies parser expectations (same framing)."""

    def _parse_cmd_as_telem(self, frame: bytes):
        """Feed a command frame through FrameParser (type codes overlap logic)."""
        results, parser = _parse_all(frame)
        return results, parser

    def test_heartbeat_roundtrip(self):
        frame = encode_heartbeat()
        results, parser = self._parse_cmd_as_telem(frame)
        assert len(results) == 1
        assert parser.frames_error == 0

    def test_speed_steer_roundtrip(self):
        frame = encode_speed_steer(750, -300)
        results, _ = self._parse_cmd_as_telem(frame)
        assert len(results) == 1
        type_code, payload = results[0]
        assert type_code == CmdType.SPEED_STEER
        speed, steer = struct.unpack(">hh", payload)
        assert speed == 750
        assert steer == -300

    def test_pid_roundtrip(self):
        frame = encode_pid_update(2.5, 0.05, 0.08)
        results, _ = self._parse_cmd_as_telem(frame)
        type_code, payload = results[0]
        kp, ki, kd = struct.unpack(">fff", payload)
        assert kp == pytest.approx(2.5, rel=1e-5)
        assert ki == pytest.approx(0.05, rel=1e-5)
        assert kd == pytest.approx(0.08, rel=1e-5)
