#!/usr/bin/env python3
"""
test_mode_switching.py — Mode transition integration tests.

Covers:
  - idle → drive: drive commands become accepted
  - drive → estop: immediate motor stop
  - MODE frame byte values match protocol constants
  - Unknown mode byte is ignored (no state change)

No ROS2 or real CAN hardware required.
Run with: python -m pytest test/test_mode_switching.py -v
"""

import struct
import pytest

from saltybot_can_e2e_test.can_mock import MockCANBus
from saltybot_can_e2e_test.protocol_defs import (
    MAMBA_CMD_VELOCITY,
    MAMBA_CMD_MODE,
    MAMBA_CMD_ESTOP,
    MODE_IDLE,
    MODE_DRIVE,
    MODE_ESTOP,
    build_mode_cmd,
    build_velocity_cmd,
    parse_velocity_cmd,
)
from saltybot_can_bridge.balance_protocol import (
    encode_velocity_cmd,
    encode_mode_cmd,
    encode_estop_cmd,
)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

class _Msg:
    def __init__(self, arb_id: int, data: bytes, is_extended_id: bool = False):
        self.arbitration_id = arb_id
        self.data = bytearray(data)
        self.is_extended_id = is_extended_id


class ModeStateMachine:
    """
    Minimal state machine tracking mode transitions and gating commands.
    """

    def __init__(self, bus: MockCANBus):
        self._bus = bus
        self._mode = MODE_IDLE

    def set_mode(self, mode: int) -> bool:
        """
        Transition to mode.  Returns True if accepted, False if invalid.
        Invalid mode values (not 0, 1, 2) are ignored.
        """
        if mode not in (MODE_IDLE, MODE_DRIVE, MODE_ESTOP):
            return False  # silently ignore

        prev_mode = self._mode
        self._mode = mode
        self._bus.send(_Msg(MAMBA_CMD_MODE, encode_mode_cmd(mode)))

        # Side-effects of entering ESTOP from DRIVE
        if mode == MODE_ESTOP and prev_mode == MODE_DRIVE:
            self._bus.send(_Msg(MAMBA_CMD_VELOCITY, encode_velocity_cmd(0.0, 0.0)))
            self._bus.send(_Msg(MAMBA_CMD_ESTOP, encode_estop_cmd(True)))

        return True

    def send_drive(self, left_mps: float, right_mps: float) -> bool:
        """
        Send a velocity command.  Returns True if forwarded, False if blocked.
        """
        if self._mode != MODE_DRIVE:
            return False
        self._bus.send(_Msg(MAMBA_CMD_VELOCITY, encode_velocity_cmd(left_mps, right_mps)))
        return True

    @property
    def mode(self) -> int:
        return self._mode


# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------

class TestIdleToDrive:
    def test_idle_to_drive_mode_frame(self, mock_can_bus):
        """Transitioning to DRIVE must emit a MODE=drive frame."""
        sm = ModeStateMachine(mock_can_bus)
        sm.set_mode(MODE_DRIVE)

        mode_frames = mock_can_bus.get_sent_frames_by_id(MAMBA_CMD_MODE)
        assert len(mode_frames) == 1
        assert bytes(mode_frames[0].data) == bytes([MODE_DRIVE])

    def test_idle_blocks_drive_commands(self, mock_can_bus):
        """In IDLE mode, drive commands must be suppressed."""
        sm = ModeStateMachine(mock_can_bus)
        # Attempt drive without entering DRIVE mode
        forwarded = sm.send_drive(1.0, 1.0)
        assert forwarded is False, "Drive cmd should be blocked in IDLE mode"

        vel_frames = mock_can_bus.get_sent_frames_by_id(MAMBA_CMD_VELOCITY)
        assert len(vel_frames) == 0

    def test_drive_mode_allows_commands(self, mock_can_bus):
        """After entering DRIVE mode, velocity commands must be forwarded."""
        sm = ModeStateMachine(mock_can_bus)
        sm.set_mode(MODE_DRIVE)
        mock_can_bus.reset()

        forwarded = sm.send_drive(0.5, 0.5)
        assert forwarded is True

        vel_frames = mock_can_bus.get_sent_frames_by_id(MAMBA_CMD_VELOCITY)
        assert len(vel_frames) == 1
        l, r = parse_velocity_cmd(bytes(vel_frames[0].data))
        assert abs(l - 0.5) < 1e-4
        assert abs(r - 0.5) < 1e-4


class TestDriveToEstop:
    def test_drive_to_estop_stops_motors(self, mock_can_bus):
        """Transitioning DRIVE → ESTOP must immediately send zero velocity."""
        sm = ModeStateMachine(mock_can_bus)
        sm.set_mode(MODE_DRIVE)
        sm.send_drive(1.0, 1.0)
        mock_can_bus.reset()

        sm.set_mode(MODE_ESTOP)

        vel_frames = mock_can_bus.get_sent_frames_by_id(MAMBA_CMD_VELOCITY)
        assert len(vel_frames) >= 1, "No velocity frame on DRIVE→ESTOP transition"
        l, r = parse_velocity_cmd(bytes(vel_frames[-1].data))
        assert abs(l) < 1e-5, f"Left motor {l} not zero after ESTOP"
        assert abs(r) < 1e-5, f"Right motor {r} not zero after ESTOP"

    def test_drive_to_estop_mode_frame(self, mock_can_bus):
        """DRIVE → ESTOP must broadcast MODE=estop."""
        sm = ModeStateMachine(mock_can_bus)
        sm.set_mode(MODE_DRIVE)
        sm.set_mode(MODE_ESTOP)

        mode_frames = mock_can_bus.get_sent_frames_by_id(MAMBA_CMD_MODE)
        assert any(bytes(f.data) == bytes([MODE_ESTOP]) for f in mode_frames)

    def test_estop_blocks_subsequent_drive(self, mock_can_bus):
        """After DRIVE → ESTOP, drive commands must be blocked."""
        sm = ModeStateMachine(mock_can_bus)
        sm.set_mode(MODE_DRIVE)
        sm.set_mode(MODE_ESTOP)
        mock_can_bus.reset()

        forwarded = sm.send_drive(1.0, 1.0)
        assert forwarded is False

        vel_frames = mock_can_bus.get_sent_frames_by_id(MAMBA_CMD_VELOCITY)
        assert len(vel_frames) == 0


class TestModeCommandEncoding:
    def test_mode_idle_byte(self, mock_can_bus):
        """MODE_IDLE must encode as 0x00."""
        assert encode_mode_cmd(MODE_IDLE) == b"\x00"

    def test_mode_drive_byte(self, mock_can_bus):
        """MODE_DRIVE must encode as 0x01."""
        assert encode_mode_cmd(MODE_DRIVE) == b"\x01"

    def test_mode_estop_byte(self, mock_can_bus):
        """MODE_ESTOP must encode as 0x02."""
        assert encode_mode_cmd(MODE_ESTOP) == b"\x02"

    def test_mode_frame_length(self, mock_can_bus):
        """Mode command payload must be exactly 1 byte."""
        for mode in (MODE_IDLE, MODE_DRIVE, MODE_ESTOP):
            payload = encode_mode_cmd(mode)
            assert len(payload) == 1, f"Mode {mode} payload length {len(payload)} != 1"

    def test_protocol_defs_build_mode_cmd_matches(self):
        """build_mode_cmd in protocol_defs must produce identical bytes."""
        for mode in (MODE_IDLE, MODE_DRIVE, MODE_ESTOP):
            assert build_mode_cmd(mode) == encode_mode_cmd(mode), \
                f"protocol_defs.build_mode_cmd({mode}) != balance_protocol.encode_mode_cmd({mode})"


class TestInvalidMode:
    def test_invalid_mode_byte_ignored(self, mock_can_bus):
        """Unknown mode byte (e.g. 0xFF) must be rejected — no state change."""
        sm = ModeStateMachine(mock_can_bus)
        sm.set_mode(MODE_DRIVE)
        initial_mode = sm.mode
        mock_can_bus.reset()

        accepted = sm.set_mode(0xFF)
        assert accepted is False, "Invalid mode 0xFF should be rejected"
        assert sm.mode == initial_mode, "Mode changed despite invalid value"
        assert len(mock_can_bus.get_sent_frames()) == 0, \
            "Frames sent for invalid mode command"

    def test_invalid_mode_99_ignored(self, mock_can_bus):
        """Mode 99 must be rejected."""
        sm = ModeStateMachine(mock_can_bus)
        accepted = sm.set_mode(99)
        assert accepted is False

    def test_invalid_mode_negative_ignored(self, mock_can_bus):
        """Negative mode values must be rejected."""
        sm = ModeStateMachine(mock_can_bus)
        accepted = sm.set_mode(-1)
        assert accepted is False

    def test_balance_protocol_invalid_mode_raises(self):
        """balance_protocol.encode_mode_cmd must raise on invalid mode."""
        with pytest.raises(ValueError):
            encode_mode_cmd(99)
        with pytest.raises(ValueError):
            encode_mode_cmd(-1)


@pytest.mark.parametrize("mode,expected_byte", [
    (MODE_IDLE,  b"\x00"),
    (MODE_DRIVE, b"\x01"),
    (MODE_ESTOP, b"\x02"),
])
def test_mode_encoding_parametrized(mode, expected_byte):
    """Parametrized check that all mode constants encode to the right byte."""
    assert encode_mode_cmd(mode) == expected_byte
