#!/usr/bin/env python3
"""
test_heartbeat_timeout.py — Tests for heartbeat loss and recovery.

Covers:
  - Heartbeat loss triggers e-stop escalation (timeout logic)
  - Heartbeat recovery restores previous mode
  - Heartbeat interval is sent at ~100 ms

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

import time
import struct
import threading
import pytest

from saltybot_can_e2e_test.can_mock import MockCANBus
from saltybot_can_e2e_test.protocol_defs import (
    ORIN_CMD_HEARTBEAT,
    ORIN_CMD_ESTOP,
    ORIN_CMD_MODE,
    MAMBA_CMD_VELOCITY,
    MAMBA_CMD_MODE,
    MAMBA_CMD_ESTOP,
    MODE_IDLE,
    MODE_DRIVE,
    MODE_ESTOP,
    build_heartbeat,
    build_estop_cmd,
    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,
)

# Heartbeat timeout from orin_can.h: 500 ms
ORIN_HB_TIMEOUT_MS = 500
ORIN_HB_TIMEOUT_S  = ORIN_HB_TIMEOUT_MS / 1000.0

# Expected heartbeat interval
HB_INTERVAL_MS = 100
HB_INTERVAL_S  = HB_INTERVAL_MS / 1000.0


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

class HeartbeatSimulator:
    """
    Simulate periodic heartbeat injection into a MockCANBus.

    Call start() to begin sending heartbeats every interval_s.
    Call stop() to cease — after ORIN_HB_TIMEOUT_S the firmware would declare
    Orin offline.
    """

    def __init__(self, bus: MockCANBus, interval_s: float = HB_INTERVAL_S):
        self._bus = bus
        self._interval_s = interval_s
        self._seq = 0
        self._running = False
        self._thread: threading.Thread | None = None

    def start(self):
        self._running = True
        self._thread = threading.Thread(target=self._run, daemon=True)
        self._thread.start()

    def stop(self):
        self._running = False

    def _run(self):
        while self._running:
            self._bus.inject(
                ORIN_CMD_HEARTBEAT,
                build_heartbeat(self._seq),
                is_extended_id=False,
            )
            self._seq += 1
            time.sleep(self._interval_s)


def _simulate_estop_on_timeout(bus: MockCANBus) -> None:
    """
    Simulate the firmware-side logic: when heartbeat timeout expires,
    the FC sends an e-stop command by setting estop mode on the ESP32-S3 BALANCE bus.
    We model this as the bridge sending zero velocity + ESTOP mode.
    """

    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, encode_velocity_cmd(0.0, 0.0)))
    bus.send(_Msg(MAMBA_CMD_MODE, encode_mode_cmd(MODE_ESTOP)))
    bus.send(_Msg(MAMBA_CMD_ESTOP, encode_estop_cmd(True)))


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

class TestHeartbeatLoss:
    def test_heartbeat_loss_triggers_estop(self, mock_can_bus):
        """
        After heartbeat ceases, the bridge must command zero velocity and
        set ESTOP mode.  We simulate this directly using _simulate_estop_on_timeout.
        """
        # First confirm the bus is clean
        assert len(mock_can_bus.get_sent_frames()) == 0

        # Simulate bridge detecting timeout and escalating
        _simulate_estop_on_timeout(mock_can_bus)

        vel_frames = mock_can_bus.get_sent_frames_by_id(MAMBA_CMD_VELOCITY)
        assert len(vel_frames) >= 1, "Zero velocity not sent after timeout"
        l, r = parse_velocity_cmd(bytes(vel_frames[-1].data))
        assert abs(l) < 1e-5, "Left not zero on timeout"
        assert abs(r) < 1e-5, "Right not zero on timeout"

        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
        ), "ESTOP mode not asserted on heartbeat timeout"

        estop_frames = mock_can_bus.get_sent_frames_by_id(MAMBA_CMD_ESTOP)
        assert len(estop_frames) >= 1, "ESTOP command not sent"
        assert bytes(estop_frames[0].data) == b"\x01"

    def test_heartbeat_loss_zero_velocity(self, mock_can_bus):
        """Zero velocity frame must appear among sent frames after timeout."""
        _simulate_estop_on_timeout(mock_can_bus)
        vel_frames = mock_can_bus.get_sent_frames_by_id(MAMBA_CMD_VELOCITY)
        assert len(vel_frames) >= 1
        for f in vel_frames:
            l, r = parse_velocity_cmd(bytes(f.data))
            assert abs(l) < 1e-5
            assert abs(r) < 1e-5


class TestHeartbeatRecovery:
    def test_heartbeat_recovery_restores_drive_mode(self, mock_can_bus):
        """
        After heartbeat loss + recovery, drive commands must be accepted again.
        We simulate: ESTOP → clear estop → send drive → verify velocity frame.
        """

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

        # Phase 1: timeout → estop
        _simulate_estop_on_timeout(mock_can_bus)
        mock_can_bus.reset()

        # Phase 2: recovery — clear estop, restore drive mode
        mock_can_bus.send(_Msg(MAMBA_CMD_ESTOP, encode_estop_cmd(False)))
        mock_can_bus.send(_Msg(MAMBA_CMD_MODE, encode_mode_cmd(MODE_DRIVE)))
        mock_can_bus.send(_Msg(MAMBA_CMD_VELOCITY, encode_velocity_cmd(0.5, 0.5)))

        estop_frames = mock_can_bus.get_sent_frames_by_id(MAMBA_CMD_ESTOP)
        assert any(bytes(f.data) == b"\x00" for f in estop_frames), \
            "ESTOP clear not sent on recovery"

        mode_frames = mock_can_bus.get_sent_frames_by_id(MAMBA_CMD_MODE)
        assert any(
            bytes(f.data) == bytes([MODE_DRIVE]) for f in mode_frames
        ), "DRIVE mode not restored after recovery"

        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[-1].data))
        assert abs(l - 0.5) < 1e-4

    def test_heartbeat_sequence_increments(self, mock_can_bus):
        """Heartbeat payloads must have incrementing sequence numbers."""
        payloads = []
        for seq in range(5):
            mock_can_bus.inject(
                ORIN_CMD_HEARTBEAT,
                build_heartbeat(seq),
                is_extended_id=False,
            )

        for i in range(5):
            frame = mock_can_bus.recv(timeout=0.05)
            assert frame is not None
            (seq_val,) = struct.unpack(">I", bytes(frame.data))
            assert seq_val == i, f"Expected seq {i}, got {seq_val}"


class TestHeartbeatInterval:
    def test_heartbeat_interval_approx_100ms(self, mock_can_bus):
        """
        Start HeartbeatSimulator, collect timestamps over ~300 ms, and verify
        the average interval is within 20% of 100 ms.
        """
        sim = HeartbeatSimulator(mock_can_bus, interval_s=0.1)
        sim.start()
        time.sleep(0.35)
        sim.stop()

        timestamps = []
        while True:
            frame = mock_can_bus.recv(timeout=0.01)
            if frame is None:
                break
            if frame.arbitration_id == ORIN_CMD_HEARTBEAT:
                timestamps.append(frame.timestamp)

        assert len(timestamps) >= 2, "Not enough heartbeat frames captured"

        intervals = [
            timestamps[i + 1] - timestamps[i]
            for i in range(len(timestamps) - 1)
        ]
        avg_interval = sum(intervals) / len(intervals)
        # Allow ±30 ms tolerance (OS scheduling jitter in CI)
        assert 0.07 <= avg_interval <= 0.13, \
            f"Average heartbeat interval {avg_interval*1000:.1f} ms not ~100 ms"

    def test_heartbeat_payload_is_4_bytes(self, mock_can_bus):
        """Heartbeat payload must be exactly 4 bytes (uint32 sequence)."""
        for seq in (0, 1, 0xFFFFFFFF):
            payload = build_heartbeat(seq)
            assert len(payload) == 4, \
                f"Heartbeat payload length {len(payload)} != 4"
