"""Tests for PID gain scheduler."""

import pytest
from std_msgs.msg import Float32
import rclpy

from saltybot_pid_scheduler.pid_scheduler_node import PIDSchedulerNode


@pytest.fixture
def rclpy_fixture():
    rclpy.init()
    yield
    rclpy.shutdown()


@pytest.fixture
def node(rclpy_fixture):
    node = PIDSchedulerNode()
    yield node
    node.destroy_node()


class TestInit:
    def test_node_initialization(self, node):
        assert node.base_kp == 1.0
        assert node.base_ki == 0.1
        assert node.base_kd == 0.05
        assert node.speed_scale == 1.0
        assert node.terrain_roughness == 0.0


class TestGainScheduling:
    def test_gains_at_baseline(self, node):
        """Test gains at baseline (no speed reduction, no terrain roughness)."""
        node._update_gains()
        
        assert node.current_kp == pytest.approx(1.0)
        assert node.current_kd == pytest.approx(0.05)

    def test_kp_increases_with_terrain(self, node):
        """Test P gain increases with terrain roughness."""
        node.terrain_roughness = 0.0
        node._update_gains()
        kp_smooth = node.current_kp

        node.terrain_roughness = 1.0
        node._update_gains()
        kp_rough = node.current_kp

        assert kp_rough > kp_smooth
        assert kp_rough == pytest.approx(1.5)  # 1.0 * (1 + 0.5 * 1.0)

    def test_kd_decreases_at_low_speed(self, node):
        """Test D gain decreases when robot slows down."""
        node.speed_scale = 1.0
        node._update_gains()
        kd_fast = node.current_kd

        node.speed_scale = 0.0  # Robot stopped
        node._update_gains()
        kd_slow = node.current_kd

        assert kd_slow < kd_fast

    def test_ki_scales_with_terrain_and_speed(self, node):
        """Test I gain scales with both terrain and speed."""
        node.speed_scale = 1.0
        node.terrain_roughness = 0.0
        node._update_gains()
        ki_baseline = node.current_ki

        node.terrain_roughness = 1.0
        node._update_gains()
        ki_rough = node.current_ki

        assert ki_rough > ki_baseline

    def test_kd_never_negative(self, node):
        """Test D gain never goes negative."""
        node.speed_scale = 0.0
        node._update_gains()

        assert node.current_kd >= 0.0

    def test_speed_scale_clipping(self, node):
        """Test speed scale is clipped to [0, 1]."""
        msg = Float32()
        msg.data = 2.0  # Out of range

        node._on_speed_scale(msg)

        assert node.speed_scale == 1.0

    def test_terrain_roughness_clipping(self, node):
        """Test terrain roughness is clipped to [0, 1]."""
        msg = Float32()
        msg.data = -0.5  # Out of range

        node._on_terrain_roughness(msg)

        assert node.terrain_roughness == 0.0


class TestSensorInputs:
    def test_speed_scale_callback(self, node):
        """Test speed scale subscription and update."""
        msg = Float32()
        msg.data = 0.5

        node._on_speed_scale(msg)

        assert node.speed_scale == 0.5

    def test_terrain_roughness_callback(self, node):
        """Test terrain roughness subscription and update."""
        msg = Float32()
        msg.data = 0.75

        node._on_terrain_roughness(msg)

        assert node.terrain_roughness == 0.75

    def test_combined_effects(self, node):
        """Test combined effect of speed and terrain on gains."""
        # Slow speed + rough terrain = high P, low D
        node.speed_scale = 0.2
        node.terrain_roughness = 0.9

        node._update_gains()

        # P should be high (due to terrain)
        assert node.current_kp > node.base_kp

        # D should be low (due to slow speed)
        # D factor = 1 + (-0.3) * (1 - 0.2) = 1 - 0.24 = 0.76
        assert node.current_kd < node.base_kd


class TestEdgeCases:
    def test_zero_speed_scale(self, node):
        """Test behavior at zero speed (robot stopped)."""
        node.speed_scale = 0.0
        node.terrain_roughness = 0.5

        node._update_gains()

        # All gains should be positive
        assert node.current_kp > 0
        assert node.current_ki > 0
        assert node.current_kd > 0

    def test_max_terrain_roughness(self, node):
        """Test behavior on extremely rough terrain."""
        node.speed_scale = 1.0
        node.terrain_roughness = 1.0

        node._update_gains()

        # Kp should be maximum
        assert node.current_kp == pytest.approx(1.5)

    def test_rapid_sensor_changes(self, node):
        """Test rapid changes in sensor inputs."""
        for speed in [1.0, 0.5, 0.1, 0.9, 1.0]:
            msg = Float32()
            msg.data = speed
            node._on_speed_scale(msg)

        for roughness in [0.0, 0.5, 1.0, 0.3, 0.0]:
            msg = Float32()
            msg.data = roughness
            node._on_terrain_roughness(msg)

        # Should end at baseline
        node._update_gains()
        assert node.speed_scale == 1.0
        assert node.terrain_roughness == 0.0
