#!/usr/bin/env python3
"""
test_estop.py — E-stop command integration tests.

Covers:
  - ESTOP command halts motors immediately
  - ESTOP persists: DRIVE commands ignored while ESTOP is active
  - ESTOP clear restores normal drive operation
  - Firmware-side estop via FC_STATUS flags is detected correctly

No ROS2 or real CAN hardware required.
Run with: python -m pytest test/test_estop.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,
    ORIN_CMD_ESTOP,
    FC_STATUS,
    MODE_IDLE,
    MODE_DRIVE,
    MODE_ESTOP,
    build_estop_cmd,
    build_mode_cmd,
    build_velocity_cmd,
    build_fc_status,
    parse_velocity_cmd,
    parse_fc_status,
)
from saltybot_can_bridge.balance_protocol import (
    encode_velocity_cmd,
    encode_mode_cmd,
    encode_estop_cmd,
)


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

class _Msg:
    """Minimal CAN message stand-in."""
    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 EstopStateMachine:
    """
    Lightweight state machine that mirrors the bridge estop logic.

    Tracks whether ESTOP is active and gates velocity commands accordingly.
    Sends frames to the supplied MockCANBus.
    """

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

    def assert_estop(self) -> None:
        """Send ESTOP and transition to estop mode."""
        self._estop_active = True
        self._mode = MODE_ESTOP
        self._bus.send(_Msg(MAMBA_CMD_VELOCITY, encode_velocity_cmd(0.0, 0.0)))
        self._bus.send(_Msg(MAMBA_CMD_MODE, encode_mode_cmd(MODE_ESTOP)))
        self._bus.send(_Msg(MAMBA_CMD_ESTOP, encode_estop_cmd(True)))

    def clear_estop(self) -> None:
        """Clear ESTOP and return to IDLE mode."""
        self._estop_active = False
        self._mode = MODE_IDLE
        self._bus.send(_Msg(MAMBA_CMD_ESTOP, encode_estop_cmd(False)))
        self._bus.send(_Msg(MAMBA_CMD_MODE, encode_mode_cmd(MODE_IDLE)))

    def send_drive(self, left_mps: float, right_mps: float) -> None:
        """Send velocity command only if ESTOP is not active."""
        if self._estop_active:
            # Bridge silently drops commands while estopped
            return
        self._mode = MODE_DRIVE
        self._bus.send(_Msg(MAMBA_CMD_VELOCITY, encode_velocity_cmd(left_mps, right_mps)))
        self._bus.send(_Msg(MAMBA_CMD_MODE, encode_mode_cmd(MODE_DRIVE)))

    @property
    def estop_active(self) -> bool:
        return self._estop_active


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

class TestEstopHaltsMotors:
    def test_estop_command_halts_motors(self, mock_can_bus):
        """
        After injecting ESTOP, zero velocity must be commanded immediately.
        """
        sm = EstopStateMachine(mock_can_bus)
        sm.assert_estop()

        vel_frames = mock_can_bus.get_sent_frames_by_id(MAMBA_CMD_VELOCITY)
        assert len(vel_frames) >= 1, "No velocity frame after ESTOP"
        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_estop_mode_frame_sent(self, mock_can_bus):
        """ESTOP mode byte must be broadcast on CAN."""
        sm = EstopStateMachine(mock_can_bus)
        sm.assert_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
        ), "MODE=ESTOP not found in sent frames"

    def test_estop_flag_byte_is_0x01(self, mock_can_bus):
        """MAMBA_CMD_ESTOP payload must be 0x01 when asserting e-stop."""
        sm = EstopStateMachine(mock_can_bus)
        sm.assert_estop()

        estop_frames = mock_can_bus.get_sent_frames_by_id(MAMBA_CMD_ESTOP)
        assert len(estop_frames) >= 1
        assert bytes(estop_frames[-1].data) == b"\x01", \
            f"ESTOP payload {estop_frames[-1].data!r} != 0x01"


class TestEstopPersists:
    def test_estop_persists_after_drive_cmd(self, mock_can_bus):
        """
        After ESTOP, injecting a DRIVE command must NOT forward velocity to the bus.
        """
        sm = EstopStateMachine(mock_can_bus)
        sm.assert_estop()
        mock_can_bus.reset()   # start fresh after initial ESTOP frames

        sm.send_drive(1.0, 1.0)   # should be suppressed

        vel_frames = mock_can_bus.get_sent_frames_by_id(MAMBA_CMD_VELOCITY)
        assert len(vel_frames) == 0, \
            "Velocity command was forwarded while ESTOP is active"

    def test_estop_mode_unchanged_after_drive_attempt(self, mock_can_bus):
        """
        After ESTOP, attempting DRIVE must not change the mode to DRIVE.
        """
        sm = EstopStateMachine(mock_can_bus)
        sm.assert_estop()
        mock_can_bus.reset()

        sm.send_drive(0.5, 0.5)

        # No mode frames should have been emitted (drive was suppressed)
        mode_frames = mock_can_bus.get_sent_frames_by_id(MAMBA_CMD_MODE)
        assert all(
            bytes(f.data) != bytes([MODE_DRIVE]) for f in mode_frames
        ), "MODE=DRIVE was set despite active ESTOP"


class TestEstopClear:
    def test_estop_clear_restores_drive(self, mock_can_bus):
        """After ESTOP_CLEAR, drive commands must be accepted again."""
        sm = EstopStateMachine(mock_can_bus)
        sm.assert_estop()
        sm.clear_estop()
        mock_can_bus.reset()

        sm.send_drive(0.8, 0.8)

        vel_frames = mock_can_bus.get_sent_frames_by_id(MAMBA_CMD_VELOCITY)
        assert len(vel_frames) == 1, "Velocity command not sent after ESTOP clear"
        l, r = parse_velocity_cmd(bytes(vel_frames[0].data))
        assert abs(l - 0.8) < 1e-4
        assert abs(r - 0.8) < 1e-4

    def test_estop_clear_flag_byte_is_0x00(self, mock_can_bus):
        """MAMBA_CMD_ESTOP payload must be 0x00 when clearing e-stop."""
        sm = EstopStateMachine(mock_can_bus)
        sm.assert_estop()
        sm.clear_estop()

        estop_frames = mock_can_bus.get_sent_frames_by_id(MAMBA_CMD_ESTOP)
        assert len(estop_frames) >= 2
        # Last ESTOP frame should be the clear
        assert bytes(estop_frames[-1].data) == b"\x00", \
            f"ESTOP clear payload {estop_frames[-1].data!r} != 0x00"

    def test_estop_clear_mode_returns_to_idle(self, mock_can_bus):
        """After clearing ESTOP, the mode frame must be MODE_IDLE."""
        sm = EstopStateMachine(mock_can_bus)
        sm.assert_estop()
        sm.clear_estop()

        mode_frames = mock_can_bus.get_sent_frames_by_id(MAMBA_CMD_MODE)
        last_mode = bytes(mode_frames[-1].data)
        assert last_mode == bytes([MODE_IDLE]), \
            f"Mode after ESTOP clear is {last_mode!r}, expected MODE_IDLE"


class TestFirmwareSideEstop:
    def test_fc_status_estop_flag_detected(self, mock_can_bus):
        """
        Simulate firmware sending estop via FC_STATUS flags (bit0=estop_active).
        Verify the Orin bridge side correctly parses the flag.
        """
        # Build FC_STATUS with estop_active bit set (flags=0x01)
        payload = build_fc_status(
            pitch_x10=0,
            motor_cmd=0,
            vbat_mv=24000,
            balance_state=2,   # TILT_FAULT
            flags=0x01,        # bit0 = estop_active
        )
        mock_can_bus.inject(FC_STATUS, payload)

        frame = mock_can_bus.recv(timeout=0.1)
        assert frame is not None, "FC_STATUS frame not received"
        parsed = parse_fc_status(bytes(frame.data))
        assert parsed["estop_active"] is True, \
            "estop_active flag not set in FC_STATUS"
        assert parsed["balance_state"] == 2

    def test_fc_status_no_estop_flag(self, mock_can_bus):
        """FC_STATUS with flags=0x00 must NOT set estop_active."""
        payload = build_fc_status(flags=0x00)
        mock_can_bus.inject(FC_STATUS, payload)
        frame = mock_can_bus.recv(timeout=0.1)
        parsed = parse_fc_status(bytes(frame.data))
        assert parsed["estop_active"] is False

    def test_fc_status_armed_flag_detected(self, mock_can_bus):
        """FC_STATUS flags bit1=armed must parse correctly."""
        payload = build_fc_status(flags=0x02)  # bit1 = armed
        mock_can_bus.inject(FC_STATUS, payload)
        frame = mock_can_bus.recv(timeout=0.1)
        parsed = parse_fc_status(bytes(frame.data))
        assert parsed["armed"] is True
        assert parsed["estop_active"] is False

    def test_fc_status_roundtrip(self, mock_can_bus):
        """build_fc_status → inject → recv → parse_fc_status must be identity."""
        payload = build_fc_status(
            pitch_x10=150,
            motor_cmd=-200,
            vbat_mv=23800,
            balance_state=1,
            flags=0x03,
        )
        mock_can_bus.inject(FC_STATUS, payload)
        frame = mock_can_bus.recv(timeout=0.1)
        parsed = parse_fc_status(bytes(frame.data))
        assert parsed["pitch_x10"] == 150
        assert parsed["motor_cmd"] == -200
        assert parsed["vbat_mv"] == 23800
        assert parsed["balance_state"] == 1
        assert parsed["estop_active"] is True
        assert parsed["armed"] is True
