#!/usr/bin/env python3
"""
test_fc_vesc_broadcast.py — FC_VESC broadcast and VESC STATUS integration tests.

Covers:
  - VESC STATUS extended frame for left VESC (ID 56) triggers FC_VESC broadcast
  - Both left (56) and right (68) VESC STATUS combined in FC_VESC
  - FC_VESC broadcast rate (~10 Hz)
  - current_x10 scaling matches protocol spec

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

import struct
import time
import threading
import pytest

from saltybot_can_e2e_test.can_mock import MockCANBus
from saltybot_can_e2e_test.protocol_defs import (
    FC_VESC,
    VESC_CAN_ID_LEFT,
    VESC_CAN_ID_RIGHT,
    VESC_STATUS_ID,
    VESC_SET_RPM_ID,
    VESC_TELEM_STATE,
    build_vesc_status,
    build_fc_vesc,
    parse_fc_vesc,
    parse_vesc_status,
)
from saltybot_can_bridge.balance_protocol import (
    VESC_TELEM_STATE as BRIDGE_VESC_TELEM_STATE,
    decode_vesc_state,
)


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

class VescStatusAggregator:
    """
    Simulates the firmware logic that:
      1. Receives VESC STATUS extended frames from left/right VESCs
      2. Builds an FC_VESC broadcast payload
      3. Injects the FC_VESC frame onto the mock bus

    This represents the ESP32-S3 BALANCE → Orin telemetry path.
    """

    def __init__(self, bus: MockCANBus):
        self._bus = bus
        self._left_rpm_x10 = 0
        self._right_rpm_x10 = 0
        self._left_current_x10 = 0
        self._right_current_x10 = 0
        self._left_seen = False
        self._right_seen = False

    def process_vesc_status(self, arb_id: int, data: bytes) -> None:
        """
        Process an incoming VESC STATUS frame (extended 29-bit ID).
        Updates internal state; broadcasts FC_VESC when at least one side is known.
        """
        node_id = arb_id & 0xFF
        parsed = parse_vesc_status(data)
        rpm_x10 = parsed["rpm"] // 10   # convert full RPM to RPM/10

        if node_id == VESC_CAN_ID_LEFT:
            self._left_rpm_x10 = rpm_x10
            self._left_current_x10 = parsed["current_x10"]
            self._left_seen = True
        elif node_id == VESC_CAN_ID_RIGHT:
            self._right_rpm_x10 = rpm_x10
            self._right_current_x10 = parsed["current_x10"]
            self._right_seen = True

        # Broadcast FC_VESC whenever we receive any update
        self._broadcast_fc_vesc()

    def _broadcast_fc_vesc(self) -> None:
        payload = build_fc_vesc(
            left_rpm_x10=self._left_rpm_x10,
            right_rpm_x10=self._right_rpm_x10,
            left_current_x10=self._left_current_x10,
            right_current_x10=self._right_current_x10,
        )
        self._bus.inject(FC_VESC, payload)


def _inject_vesc_status(bus: MockCANBus, vesc_id: int, rpm: int,
                        current_x10: int = 50, duty_x1000: int = 250) -> None:
    """Inject a VESC STATUS extended frame for the given node ID."""
    arb_id = VESC_STATUS_ID(vesc_id)
    payload = build_vesc_status(rpm=rpm, current_x10=current_x10, duty_x1000=duty_x1000)
    bus.inject(arb_id, payload, is_extended_id=True)


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

class TestVescStatusToFcVesc:
    def test_left_vesc_status_triggers_broadcast(self, mock_can_bus):
        """
        Inject VESC STATUS for left VESC (ID 56) → verify FC_VESC contains
        the correct left RPM (rpm / 10).
        """
        agg = VescStatusAggregator(mock_can_bus)

        # Left VESC: 3000 RPM → rpm_x10 = 300
        arb_id = VESC_STATUS_ID(VESC_CAN_ID_LEFT)
        payload = build_vesc_status(rpm=3000, current_x10=55)
        agg.process_vesc_status(arb_id, payload)

        frame = mock_can_bus.recv(timeout=0.1)
        assert frame is not None, "No FC_VESC broadcast after left VESC STATUS"
        parsed = parse_fc_vesc(bytes(frame.data))
        assert parsed["left_rpm_x10"] == 300, \
            f"left_rpm_x10 {parsed['left_rpm_x10']} != 300"
        assert abs(parsed["left_rpm"] - 3000.0) < 1.0

    def test_right_vesc_status_triggers_broadcast(self, mock_can_bus):
        """Inject VESC STATUS for right VESC (ID 68) → verify right RPM in FC_VESC."""
        agg = VescStatusAggregator(mock_can_bus)

        arb_id = VESC_STATUS_ID(VESC_CAN_ID_RIGHT)
        payload = build_vesc_status(rpm=2000, current_x10=40)
        agg.process_vesc_status(arb_id, payload)

        frame = mock_can_bus.recv(timeout=0.1)
        assert frame is not None
        parsed = parse_fc_vesc(bytes(frame.data))
        assert parsed["right_rpm_x10"] == 200

    def test_left_vesc_id_matches_constant(self):
        """VESC_STATUS_ID(56) must equal (9 << 8) | 56 = 0x938."""
        assert VESC_STATUS_ID(VESC_CAN_ID_LEFT) == (9 << 8) | 56
        assert VESC_STATUS_ID(VESC_CAN_ID_LEFT) == 0x938

    def test_right_vesc_id_matches_constant(self):
        """VESC_STATUS_ID(68) must equal (9 << 8) | 68 = 0x944."""
        assert VESC_STATUS_ID(VESC_CAN_ID_RIGHT) == (9 << 8) | 68
        assert VESC_STATUS_ID(VESC_CAN_ID_RIGHT) == 0x944


class TestBothVescStatusCombined:
    def test_both_vesc_status_combined_in_fc_vesc(self, mock_can_bus):
        """
        Inject both left (56) and right (68) VESC STATUS frames.
        Final FC_VESC must contain both RPMs.
        """
        agg = VescStatusAggregator(mock_can_bus)

        # Left: 3000 RPM
        agg.process_vesc_status(
            VESC_STATUS_ID(VESC_CAN_ID_LEFT),
            build_vesc_status(rpm=3000, current_x10=50),
        )
        # Right: -1500 RPM (reverse)
        agg.process_vesc_status(
            VESC_STATUS_ID(VESC_CAN_ID_RIGHT),
            build_vesc_status(rpm=-1500, current_x10=30),
        )

        # Drain two FC_VESC frames (one per update), check the latest
        frames = []
        while True:
            f = mock_can_bus.recv(timeout=0.05)
            if f is None:
                break
            frames.append(f)

        assert len(frames) >= 2, "Expected at least 2 FC_VESC frames"
        # Last frame must have both sides
        last = parse_fc_vesc(bytes(frames[-1].data))
        assert last["left_rpm_x10"] == 300, \
            f"left_rpm_x10 {last['left_rpm_x10']} != 300"
        assert last["right_rpm_x10"] == -150, \
            f"right_rpm_x10 {last['right_rpm_x10']} != -150"

    def test_both_vesc_currents_combined(self, mock_can_bus):
        """Both current values must appear in FC_VESC after two STATUS frames."""
        agg = VescStatusAggregator(mock_can_bus)
        agg.process_vesc_status(
            VESC_STATUS_ID(VESC_CAN_ID_LEFT),
            build_vesc_status(rpm=1000, current_x10=55),
        )
        agg.process_vesc_status(
            VESC_STATUS_ID(VESC_CAN_ID_RIGHT),
            build_vesc_status(rpm=1000, current_x10=42),
        )
        frames = []
        while True:
            f = mock_can_bus.recv(timeout=0.05)
            if f is None:
                break
            frames.append(f)
        last = parse_fc_vesc(bytes(frames[-1].data))
        assert last["left_current_x10"] == 55
        assert last["right_current_x10"] == 42


class TestVescBroadcastRate:
    def test_fc_vesc_broadcast_at_10hz(self, mock_can_bus):
        """
        Simulate FC_VESC broadcasts at ~10 Hz and verify the rate.
        We inject 12 frames over ~120 ms, then verify count and average interval.
        """
        _FC_VESC_HZ = 10
        _interval = 1.0 / _FC_VESC_HZ

        timestamps = []
        stop_event = threading.Event()

        def broadcaster():
            while not stop_event.is_set():
                t = time.monotonic()
                mock_can_bus.inject(
                    FC_VESC,
                    build_fc_vesc(100, -100, 30, 30),
                    timestamp=t,
                )
                timestamps.append(t)
                time.sleep(_interval)

        t = threading.Thread(target=broadcaster, daemon=True)
        t.start()
        time.sleep(0.15)   # collect ~1.5 broadcasts
        stop_event.set()
        t.join(timeout=0.2)

        assert len(timestamps) >= 1, "No FC_VESC broadcasts in 150 ms window"

        if len(timestamps) >= 2:
            intervals = [timestamps[i+1] - timestamps[i] for i in range(len(timestamps)-1)]
            avg = sum(intervals) / len(intervals)
            # ±40 ms tolerance for OS scheduling
            assert 0.06 <= avg <= 0.14, \
                f"FC_VESC broadcast interval {avg*1000:.1f} ms not ~100 ms"

    def test_fc_vesc_frame_is_8_bytes(self):
        """FC_VESC payload must always be exactly 8 bytes."""
        payload = build_fc_vesc(300, -150, 55, 42)
        assert len(payload) == 8


class TestVescCurrentScaling:
    def test_current_x10_scaling(self, mock_can_bus):
        """
        Verify current_x10 scaling: 5.5 A → current_x10=55.
        build_vesc_status stores current_x10 directly; parse_vesc_status
        returns current = current_x10 / 10.
        """
        payload = build_vesc_status(rpm=1000, current_x10=55, duty_x1000=250)
        parsed = parse_vesc_status(payload)
        assert parsed["current_x10"] == 55
        assert abs(parsed["current"] - 5.5) < 0.01

    def test_current_negative_scaling(self, mock_can_bus):
        """Negative current (regen) must scale correctly."""
        payload = build_vesc_status(rpm=-500, current_x10=-30)
        parsed = parse_vesc_status(payload)
        assert parsed["current_x10"] == -30
        assert abs(parsed["current"] - (-3.0)) < 0.01

    def test_fc_vesc_current_x10_roundtrip(self, mock_can_bus):
        """build_fc_vesc → inject → recv → parse must preserve current_x10."""
        payload = build_fc_vesc(
            left_rpm_x10=200,
            right_rpm_x10=200,
            left_current_x10=55,
            right_current_x10=42,
        )
        mock_can_bus.inject(FC_VESC, payload)
        frame = mock_can_bus.recv(timeout=0.1)
        parsed = parse_fc_vesc(bytes(frame.data))
        assert parsed["left_current_x10"] == 55
        assert parsed["right_current_x10"] == 42

    @pytest.mark.parametrize("vesc_id", [VESC_CAN_ID_LEFT, VESC_CAN_ID_RIGHT])
    def test_vesc_status_id_both_nodes(self, vesc_id, mock_can_bus):
        """
        VESC_STATUS_ID(vesc_id) must produce the correct 29-bit extended arb_id
        for both the left (56) and right (68) node IDs.
        """
        expected = (9 << 8) | vesc_id
        assert VESC_STATUS_ID(vesc_id) == expected

    @pytest.mark.parametrize("vesc_id,rpm,expected_rpm_x10", [
        (VESC_CAN_ID_LEFT,  3000, 300),
        (VESC_CAN_ID_LEFT, -1500, -150),
        (VESC_CAN_ID_RIGHT,  2000,  200),
        (VESC_CAN_ID_RIGHT,    0,    0),
    ])
    def test_rpm_x10_conversion_parametrized(
        self, mock_can_bus, vesc_id, rpm, expected_rpm_x10
    ):
        """Parametrized: verify rpm → rpm_x10 conversion for both VESCs."""
        agg = VescStatusAggregator(mock_can_bus)
        agg.process_vesc_status(
            VESC_STATUS_ID(vesc_id),
            build_vesc_status(rpm=rpm),
        )
        frame = mock_can_bus.recv(timeout=0.05)
        assert frame is not None
        parsed = parse_fc_vesc(bytes(frame.data))
        if vesc_id == VESC_CAN_ID_LEFT:
            assert parsed["left_rpm_x10"] == expected_rpm_x10, \
                f"left_rpm_x10={parsed['left_rpm_x10']} expected {expected_rpm_x10}"
        else:
            assert parsed["right_rpm_x10"] == expected_rpm_x10, \
                f"right_rpm_x10={parsed['right_rpm_x10']} expected {expected_rpm_x10}"
