#!/usr/bin/env python3
"""
test_drive_command.py — Integration tests for the drive command path.

Tests verify:
  DRIVE cmd → ESP32-S3 BALANCE receives velocity command frame → mock VESC status response
  → FC_VESC broadcast contains correct RPMs.

All tests run without real hardware or a running ROS2 system.
Run with: python -m pytest test/test_drive_command.py -v
"""

import struct
import pytest

from saltybot_can_e2e_test.protocol_defs import (
    MAMBA_CMD_VELOCITY,
    MAMBA_CMD_MODE,
    FC_VESC,
    MODE_DRIVE,
    MODE_IDLE,
    VESC_CAN_ID_LEFT,
    VESC_CAN_ID_RIGHT,
    VESC_STATUS_ID,
    build_velocity_cmd,
    build_fc_vesc,
    build_vesc_status,
    parse_velocity_cmd,
    parse_fc_vesc,
)
from saltybot_can_bridge.balance_protocol import (
    encode_velocity_cmd,
    encode_mode_cmd,
)


# ---------------------------------------------------------------------------
# Helper
# ---------------------------------------------------------------------------

def _send_drive(bus, left_mps: float, right_mps: float) -> None:
    """Simulate the bridge encoding and sending a velocity command."""
    from saltybot_can_e2e_test.can_mock import MockCANBus

    payload = encode_velocity_cmd(left_mps, right_mps)
    # Create a minimal message object compatible with our mock
    class _Msg:
        def __init__(self, arb_id, data):
            self.arbitration_id = arb_id
            self.data = bytearray(data)
            self.is_extended_id = False

    bus.send(_Msg(MAMBA_CMD_VELOCITY, payload))
    bus.send(_Msg(MAMBA_CMD_MODE, encode_mode_cmd(MODE_DRIVE)))


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

class TestDriveForward:
    def test_drive_forward_velocity_frame_sent(self, mock_can_bus):
        """
        Inject DRIVE cmd (1.0 m/s, 1.0 m/s) → verify ESP32-S3 BALANCE receives
        a MAMBA_CMD_VELOCITY frame with correct payload.
        """
        _send_drive(mock_can_bus, 1.0, 1.0)

        vel_frames = mock_can_bus.get_sent_frames_by_id(MAMBA_CMD_VELOCITY)
        assert len(vel_frames) == 1, "Expected exactly one velocity command frame"

        left, right = parse_velocity_cmd(bytes(vel_frames[0].data))
        assert abs(left - 1.0) < 1e-4, f"Left speed {left} != 1.0"
        assert abs(right - 1.0) < 1e-4, f"Right speed {right} != 1.0"

    def test_drive_forward_mode_drive_sent(self, mock_can_bus):
        """After a drive command, a MODE=drive frame must accompany it."""
        _send_drive(mock_can_bus, 1.0, 1.0)

        mode_frames = mock_can_bus.get_sent_frames_by_id(MAMBA_CMD_MODE)
        assert len(mode_frames) >= 1, "Expected at least one MODE frame"
        assert bytes(mode_frames[0].data) == bytes([MODE_DRIVE])

    def test_drive_forward_fc_vesc_broadcast(self, mock_can_bus):
        """
        Simulate FC_VESC broadcast arriving after drive cmd; verify parse is correct.
        (In the real loop ESP32-S3 BALANCE computes RPM from m/s and broadcasts FC_VESC.)
        This test checks the FC_VESC frame format and parser.
        """
        # Simulate: 1.0 m/s → ~300 RPM × 10 = 3000 (representative, not physics)
        fc_payload = build_fc_vesc(
            left_rpm_x10=300, right_rpm_x10=300,
            left_current_x10=50, right_current_x10=50,
        )
        mock_can_bus.inject(FC_VESC, fc_payload)

        frame = mock_can_bus.recv(timeout=0.1)
        assert frame is not None, "FC_VESC frame not received"
        parsed = parse_fc_vesc(bytes(frame.data))
        assert parsed["left_rpm_x10"] == 300
        assert parsed["right_rpm_x10"] == 300
        assert abs(parsed["left_rpm"] - 3000.0) < 0.1


class TestDriveTurn:
    def test_drive_turn_differential_rpm(self, mock_can_bus):
        """
        DRIVE cmd (0.5, −0.5) → verify differential RPM in velocity command.
        """
        _send_drive(mock_can_bus, 0.5, -0.5)

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

        left, right = parse_velocity_cmd(bytes(vel_frames[0].data))
        assert abs(left - 0.5) < 1e-4,  f"Left speed {left} != 0.5"
        assert abs(right - (-0.5)) < 1e-4, f"Right speed {right} != -0.5"
        # Signs must be opposite for a zero-radius spin
        assert left > 0 and right < 0

    def test_drive_turn_fc_vesc_differential(self, mock_can_bus):
        """Simulated FC_VESC for a turn has opposite-sign RPMs."""
        fc_payload = build_fc_vesc(
            left_rpm_x10=150, right_rpm_x10=-150,
            left_current_x10=30, right_current_x10=30,
        )
        mock_can_bus.inject(FC_VESC, fc_payload)
        frame = mock_can_bus.recv(timeout=0.1)
        parsed = parse_fc_vesc(bytes(frame.data))
        assert parsed["left_rpm_x10"] > 0
        assert parsed["right_rpm_x10"] < 0


class TestDriveZero:
    def test_drive_zero_stops_motors(self, mock_can_bus):
        """
        After a non-zero drive cmd, sending zero velocity must result in
        RPM=0 being commanded to both VESCs.
        """
        _send_drive(mock_can_bus, 1.0, 1.0)
        mock_can_bus.reset()          # clear prior frames

        _send_drive(mock_can_bus, 0.0, 0.0)

        vel_frames = mock_can_bus.get_sent_frames_by_id(MAMBA_CMD_VELOCITY)
        assert len(vel_frames) == 1
        left, right = parse_velocity_cmd(bytes(vel_frames[0].data))
        assert abs(left) < 1e-5,  "Left motor not stopped"
        assert abs(right) < 1e-5, "Right motor not stopped"


class TestDriveCmdTimeout:
    def test_drive_cmd_timeout_sends_zero(self, mock_can_bus):
        """
        Simulate the watchdog behaviour: if no DRIVE cmd arrives for >500 ms,
        zero velocity is sent.  We test the encoding directly (without timers).
        """
        # The watchdog in CanBridgeNode calls encode_velocity_cmd(0.0, 0.0) and
        # sends it on MAMBA_CMD_VELOCITY.  Replicate that here.
        zero_payload = encode_velocity_cmd(0.0, 0.0)

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

        mock_can_bus.send(_Msg(MAMBA_CMD_VELOCITY, zero_payload))
        mock_can_bus.send(_Msg(MAMBA_CMD_MODE, encode_mode_cmd(MODE_IDLE)))

        vel_frames = mock_can_bus.get_sent_frames_by_id(MAMBA_CMD_VELOCITY)
        assert len(vel_frames) == 1
        left, right = parse_velocity_cmd(bytes(vel_frames[0].data))
        assert abs(left) < 1e-5
        assert abs(right) < 1e-5

        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_IDLE])


@pytest.mark.parametrize("left_mps,right_mps", [
    (0.5, 0.5),
    (1.0, 0.0),
    (0.0, -1.0),
    (-0.5, -0.5),
])
def test_drive_cmd_payload_roundtrip(mock_can_bus, left_mps, right_mps):
    """Parametrized: encode then decode must recover original velocities."""
    payload = encode_velocity_cmd(left_mps, right_mps)
    l, r = parse_velocity_cmd(payload)
    assert abs(l - left_mps) < 1e-4
    assert abs(r - right_mps) < 1e-4
