import numpy as np
import pytest

from k1link.simulation.ai_polygon.contracts import WorldSettings
from k1link.simulation.ai_polygon.mission_policy import WaypointMission, inclination


def pose(x=0, y=0):
    return [x, y, 0.37, 0, 0, 0, 1]


def goal(target, prior, excluded):
    return [2, 0.7 * len(excluded), 0]


def test_chassis_jitter_cannot_hide_stall_and_retries_are_bounded():
    mission = WaypointMission([[8, 0]])
    states = []
    for t in np.arange(0, 34, 0.2):
        _, intent = mission.update(pose(0.01 * np.sin(t * 10)), float(t), goal)
        states.append(intent["state"])
    assert "replanning" in states
    assert states[-1] == "stuck"
    assert mission.update(pose(1), 35, goal)[1]["state"] == "stuck"


def test_route_cursor_and_completion_survive_pause_without_restarting_task():
    mission = WaypointMission([[1, 0], [3, 0]])
    mission.update(pose(), 0, goal)
    assert mission.update(pose(1), 6, goal)[1]["waypoint"] == 1
    mission.resume()
    assert mission.update(pose(1), 6, goal)[1]["waypoint"] == 1
    assert mission.update(pose(3), 18, goal)[1]["state"] == "goal-reached"
    mission.resume()
    assert mission.update(pose(3), 18, goal)[1]["state"] == "goal-reached"


def test_overturned_and_excessive_tilt_are_latched_before_goal_success():
    mission = WaypointMission([[0, 0]])
    overturned = [0, 0, 0.37, 1, 0, 0, 0]
    assert inclination(overturned) == pytest.approx(180)
    assert mission.update(overturned, 0, goal)[1]["state"] == "unstable"
    mission.resume()
    assert mission.update(pose(), 1, goal)[1]["state"] == "unstable"


def test_missing_surface_never_becomes_a_drive_permission_or_unbounded_wait():
    mission = WaypointMission([[8, 0]])
    for t in range(34):
        target, intent = mission.update(pose(), t, lambda *_: None)
        assert target is None
    assert intent["state"] == "stuck"


def test_route_contract_bounds_and_finiteness():
    from pydantic import ValidationError

    for route in ([[float("nan"), 0]], [[0, 0]] * 33, [[0, 10001]]):
        with pytest.raises(ValidationError):
            WorldSettings(route_xy=route)


def test_recovery_uses_observed_goal_and_does_not_count_retreat_as_progress():
    mission = WaypointMission([[8, 0]])

    def retreat(prior):
        return prior or [-0.65, 0, 0]

    mission.update(pose(), 0, goal, retreat)
    target, intent = mission.update(pose(), 8, goal, retreat)
    assert target == [-0.65, 0, 0] and intent["state"] == "reversing"
    assert mission.update(pose(-0.2), 10, goal, retreat)[1]["state"] == "reversing"
    assert mission.update(pose(-0.4), 12, goal, retreat)[1]["state"] == "replanning"
    mission.update(pose(), 16, goal, retreat)
    assert mission.attempts == 1
    # Moving back to where we started cannot create a fresh three-attempt budget.
    assert mission.update(pose(), 20, goal, retreat)[1]["recovery_attempt"] == 2


def test_recovery_stops_immediately_on_lost_support_and_times_out_without_motion():
    mission = WaypointMission([[8, 0]])

    def retreat(prior):
        return prior or [-0.65, 0, 0]

    mission.update(pose(), 0, goal, retreat)
    mission.update(pose(), 8, goal, retreat)
    assert mission.update(pose(), 8.2, goal, lambda _: None)[0] is None
    assert mission.recovery_goal is None
    for seconds in (17, 23, 32, 38, 47):
        _, intent = mission.update(pose(), seconds, goal, retreat)
    assert intent["state"] == "stuck"


def test_wrong_direction_never_replenishes_route_attempts():
    mission = WaypointMission([[8, 0]])
    for seconds in range(34):
        _, intent = mission.update(pose(-seconds * 0.1), seconds, goal)
    assert intent["state"] == "stuck"


def test_signed_motion_contract_accepts_reverse_but_stays_bounded():
    from pydantic import ValidationError

    from k1link.simulation.ai_polygon.contracts import Decision

    assert (
        Decision(
            speed_mps=-0.1, yaw_rate_rps=0, reason="replanning", road_fraction=0, obstacle_count=0
        ).speed_mps
        < 0
    )
    with pytest.raises(ValidationError):
        Decision(speed_mps=-1.1, yaw_rate_rps=0, reason="road", road_fraction=1, obstacle_count=0)


def test_slow_regulated_progress_is_not_mistaken_for_stall():
    mission = WaypointMission([[8, 0]])
    for seconds in range(60):
        _, intent = mission.update(pose(seconds * 0.02), seconds, goal)
        assert intent["state"] == "following"
        assert intent["recovery_attempt"] == 0


def test_rejected_observation_stops_but_does_not_forget_revalidated_goal():
    mission = WaypointMission([[2, 0]])
    observed = [2, 0, 0]
    assert mission.update(pose(), 0, lambda *_: observed)[0] == observed
    assert mission.update(pose(1), 1, lambda *_: None)[0] is None
    assert mission.goal == observed
    mission.resume()
    calls = []

    def revalidate(_target, prior, _excluded):
        calls.append(prior)
        return prior

    assert mission.update(pose(1.3), 1.2, revalidate)[0] == observed
    assert calls == [observed]
    # Memory alone never authorizes a command when the new frame is rejected.
    assert mission.update(pose(1.3), 1.4, lambda *_: None)[0] is None