#!/usr/bin/env python3
"""Unit tests for saltybot_esp32_serial.esp32_balance_protocol.

No ROS2 or hardware required — exercises encode/decode round-trips,
framing, CRC8, and parser state-machine in pure Python.

Run with:  pytest test/ -v
"""

import struct
import unittest

from saltybot_esp32_serial.esp32_balance_protocol import (
    SYNC_BYTE,
    CmdType,
    TelType,
    BalanceState,
    StatusFrame,
    VescFrame,
    AckFrame,
    NackFrame,
    FrameParser,
    _crc8,
    encode_heartbeat,
    encode_drive,
    encode_estop,
    encode_arm,
    encode_pid,
)


# ── CRC8 tests ────────────────────────────────────────────────────────────────

class TestCrc8(unittest.TestCase):
    """Verify CRC8-SMBUS implementation."""

    def test_empty(self):
        self.assertEqual(_crc8(b""), 0x00)

    def test_single_byte_zero(self):
        # CRC of 0x00: no XOR, no poly — stays 0
        self.assertEqual(_crc8(b"\x00"), 0x00)

    def test_single_byte_nonzero(self):
        # Pre-computed expected value for 0xFF
        result = _crc8(b"\xFF")
        self.assertIsInstance(result, int)
        self.assertGreaterEqual(result, 0)
        self.assertLessEqual(result, 0xFF)

    def test_deterministic(self):
        data = b"\x01\x02\x04\x08"
        self.assertEqual(_crc8(data), _crc8(data))

    def test_different_data_different_crc(self):
        self.assertNotEqual(_crc8(b"\x01"), _crc8(b"\x02"))

    def test_byte_range(self):
        for b in range(256):
            r = _crc8(bytes([b]))
            self.assertGreaterEqual(r, 0)
            self.assertLessEqual(r, 255)


# ── Frame structure tests ─────────────────────────────────────────────────────

class TestFrameStructure(unittest.TestCase):
    """Verify encode functions produce correctly structured frames."""

    def _check_frame(self, data: bytes, expected_type: int, expected_payload: bytes):
        """Assert frame matches [SYNC][LEN][TYPE][PAYLOAD][CRC8]."""
        self.assertGreaterEqual(len(data), 4)
        self.assertEqual(data[0], SYNC_BYTE, "SYNC byte mismatch")
        length = data[1]
        self.assertEqual(data[2], expected_type, "TYPE byte mismatch")
        self.assertEqual(length, len(expected_payload), "LEN byte mismatch")
        self.assertEqual(data[3:3 + length], expected_payload, "PAYLOAD mismatch")
        # Verify CRC over [LEN][TYPE][PAYLOAD]
        crc_data = bytes([length, expected_type]) + expected_payload
        expected_crc = _crc8(crc_data)
        self.assertEqual(data[-1], expected_crc, "CRC8 mismatch")

    def test_heartbeat_frame(self):
        frame = encode_heartbeat()
        self._check_frame(frame, CmdType.HEARTBEAT, b"")

    def test_drive_frame_zeros(self):
        frame = encode_drive(0, 0)
        payload = struct.pack(">hh", 0, 0)
        self._check_frame(frame, CmdType.DRIVE, payload)

    def test_drive_frame_values(self):
        frame = encode_drive(500, -300)
        payload = struct.pack(">hh", 500, -300)
        self._check_frame(frame, CmdType.DRIVE, payload)

    def test_drive_clamping(self):
        # Values outside ±1000 must be clamped
        frame = encode_drive(5000, -9999)
        _, _, _, p0, p1, p2, p3, _ = frame
        speed, steer = struct.unpack(">hh", bytes([p0, p1, p2, p3]))
        self.assertEqual(speed, 1000)
        self.assertEqual(steer, -1000)

    def test_estop_assert(self):
        frame = encode_estop(True)
        self._check_frame(frame, CmdType.ESTOP, b"\x01")

    def test_estop_clear(self):
        frame = encode_estop(False)
        self._check_frame(frame, CmdType.ESTOP, b"\x00")

    def test_arm_frame(self):
        for arm_val in (True, False):
            frame = encode_arm(arm_val)
            self._check_frame(frame, CmdType.ARM, bytes([1 if arm_val else 0]))

    def test_pid_frame(self):
        frame = encode_pid(1.0, 0.1, 0.01)
        payload = struct.pack(">fff", 1.0, 0.1, 0.01)
        self._check_frame(frame, CmdType.PID, payload)

    def test_pid_negative_raises(self):
        with self.assertRaises(ValueError):
            encode_pid(-1.0, 0.1, 0.01)


# ── Parser round-trip tests ───────────────────────────────────────────────────

class TestFrameParser(unittest.TestCase):
    """Feed encoded frames back through the parser and verify decode."""

    def _parse_all(self, data: bytes):
        """Feed all bytes; return list of decoded frames."""
        parser = FrameParser()
        results = []
        for byte in data:
            r = parser.feed(byte)
            if r is not None:
                results.append(r)
        return parser, results

    # ── Command echo round-trips (parser recognises unknown cmd types as raw tuples) ──

    def test_heartbeat_roundtrip(self):
        frame = encode_heartbeat()
        parser, results = self._parse_all(frame)
        self.assertEqual(len(results), 1)
        # HEARTBEAT (0x01) is a command type — no dataclass decoder, returns raw tuple
        self.assertIsInstance(results[0], tuple)
        type_code, payload = results[0]
        self.assertEqual(type_code, CmdType.HEARTBEAT)
        self.assertEqual(payload, b"")

    def test_drive_roundtrip(self):
        frame = encode_drive(250, -100)
        parser, results = self._parse_all(frame)
        self.assertEqual(len(results), 1)
        type_code, payload = results[0]
        self.assertEqual(type_code, CmdType.DRIVE)
        speed, steer = struct.unpack(">hh", payload)
        self.assertEqual(speed, 250)
        self.assertEqual(steer, -100)

    # ── Telemetry decode round-trips ────────────────────────────────────────

    def _build_telem_frame(self, type_code: int, payload: bytes) -> bytes:
        length = len(payload)
        crc = _crc8(bytes([length, type_code]) + payload)
        return bytes([SYNC_BYTE, length, type_code]) + payload + bytes([crc])

    def test_status_frame_decode(self):
        raw = struct.pack(">hhHBB", 152, 300, 11400, BalanceState.ARMED, 0x00)
        frame = self._build_telem_frame(TelType.STATUS, raw)
        _, results = self._parse_all(frame)
        self.assertEqual(len(results), 1)
        f = results[0]
        self.assertIsInstance(f, StatusFrame)
        self.assertAlmostEqual(f.pitch_deg, 15.2, places=5)
        self.assertEqual(f.motor_cmd, 300)
        self.assertEqual(f.vbat_mv, 11400)
        self.assertEqual(f.balance_state, BalanceState.ARMED)
        self.assertEqual(f.flags, 0x00)

    def test_status_estop_flag(self):
        raw = struct.pack(">hhHBB", 0, 0, 10000, BalanceState.ESTOP, 0x01)
        frame = self._build_telem_frame(TelType.STATUS, raw)
        _, results = self._parse_all(frame)
        f = results[0]
        self.assertIsInstance(f, StatusFrame)
        self.assertTrue(bool(f.flags & 0x01))  # estop_active flag set

    def test_vesc_left_frame_decode(self):
        # erpm=12000, voltage_mv=22400, current_ma=4500, temp_x10=352
        raw = struct.pack(">iHhH", 12000, 22400, 4500, 352)
        frame = self._build_telem_frame(TelType.VESC_LEFT, raw)
        _, results = self._parse_all(frame)
        self.assertEqual(len(results), 1)
        f = results[0]
        self.assertIsInstance(f, VescFrame)
        self.assertEqual(f.vesc_id, 56)
        self.assertEqual(f.erpm, 12000)
        self.assertEqual(f.voltage_mv, 22400)
        self.assertEqual(f.current_ma, 4500)
        self.assertAlmostEqual(f.temp_c, 35.2, places=5)

    def test_vesc_right_frame_decode(self):
        raw = struct.pack(">iHhH", -6000, 22000, -200, 280)
        frame = self._build_telem_frame(TelType.VESC_RIGHT, raw)
        _, results = self._parse_all(frame)
        f = results[0]
        self.assertIsInstance(f, VescFrame)
        self.assertEqual(f.vesc_id, 68)
        self.assertEqual(f.erpm, -6000)
        self.assertAlmostEqual(f.temp_c, 28.0, places=5)

    def test_ack_frame_decode(self):
        raw = bytes([CmdType.DRIVE])
        frame = self._build_telem_frame(TelType.ACK, raw)
        _, results = self._parse_all(frame)
        f = results[0]
        self.assertIsInstance(f, AckFrame)
        self.assertEqual(f.echoed_type, CmdType.DRIVE)

    def test_nack_frame_decode(self):
        raw = bytes([CmdType.DRIVE, 0x03])  # ERR_ESTOP_ACTIVE
        frame = self._build_telem_frame(TelType.NACK, raw)
        _, results = self._parse_all(frame)
        f = results[0]
        self.assertIsInstance(f, NackFrame)
        self.assertEqual(f.cmd_type, CmdType.DRIVE)
        self.assertEqual(f.error_code, 0x03)

    # ── Parser robustness ───────────────────────────────────────────────────

    def test_bad_crc_discarded(self):
        frame = bytearray(encode_heartbeat())
        frame[-1] ^= 0xFF          # corrupt CRC
        parser, results = self._parse_all(bytes(frame))
        self.assertEqual(results, [])
        self.assertEqual(parser.frames_error, 1)
        self.assertEqual(parser.frames_ok, 0)

    def test_garbage_before_sync(self):
        garbage = bytes([0x00, 0xFF, 0x01, 0x22])
        frame = encode_heartbeat()
        parser, results = self._parse_all(garbage + frame)
        self.assertEqual(len(results), 1)

    def test_two_frames_sequential(self):
        frames = encode_heartbeat() + encode_drive(100, 50)
        parser, results = self._parse_all(frames)
        self.assertEqual(len(results), 2)
        self.assertEqual(parser.frames_ok, 2)

    def test_oversized_payload_rejected(self):
        # Craft a frame claiming LEN=63 (> MAX_PAYLOAD_LEN=62)
        bad = bytes([SYNC_BYTE, 63, CmdType.DRIVE])
        parser, results = self._parse_all(bad)
        self.assertEqual(results, [])
        self.assertEqual(parser.frames_error, 1)

    def test_parser_counter_increments(self):
        parser = FrameParser()
        frame = encode_drive(10, 10)
        for b in frame:
            parser.feed(b)
        self.assertEqual(parser.frames_ok, 1)
        self.assertEqual(parser.frames_error, 0)

    def test_reset_clears_state(self):
        parser = FrameParser()
        # Feed partial frame then reset
        partial = encode_drive(0, 0)[:-2]
        for b in partial:
            parser.feed(b)
        parser.reset()
        # Now a clean frame should decode correctly
        for b in encode_heartbeat():
            r = parser.feed(b)
        self.assertEqual(parser.frames_ok, 1)


# ── Encode parameter edge cases ───────────────────────────────────────────────

class TestEncodeEdgeCases(unittest.TestCase):
    def test_drive_clamp_positive(self):
        frame = encode_drive(9999, 9999)
        payload = frame[3:7]
        speed, steer = struct.unpack(">hh", payload)
        self.assertEqual(speed, 1000)
        self.assertEqual(steer, 1000)

    def test_drive_clamp_negative(self):
        frame = encode_drive(-9999, -9999)
        payload = frame[3:7]
        speed, steer = struct.unpack(">hh", payload)
        self.assertEqual(speed, -1000)
        self.assertEqual(steer, -1000)

    def test_pid_zero_gains_valid(self):
        frame = encode_pid(0.0, 0.0, 0.0)
        self.assertIsNotNone(frame)

    def test_pid_large_gains_valid(self):
        frame = encode_pid(100.0, 50.0, 25.0)
        self.assertIsNotNone(frame)


if __name__ == "__main__":
    unittest.main()
