from __future__ import annotations

from env.environment import LastMileDeliveryEnvironment
from env.models import ActionType, Direction, EnvironmentConfig, OrderPriority, SimulatorAction


def _move_towards(env: LastMileDeliveryEnvironment, target_x: int, target_y: int) -> tuple[float, bool]:
    obs = env.current_observation()
    if obs.agent_location.x < target_x:
        _, reward, done, _ = env.step(SimulatorAction(action_type=ActionType.MOVE, direction=Direction.RIGHT))
        return reward, done
    if obs.agent_location.x > target_x:
        _, reward, done, _ = env.step(SimulatorAction(action_type=ActionType.MOVE, direction=Direction.LEFT))
        return reward, done
    if obs.agent_location.y < target_y:
        _, reward, done, _ = env.step(SimulatorAction(action_type=ActionType.MOVE, direction=Direction.DOWN))
        return reward, done
    _, reward, done, _ = env.step(SimulatorAction(action_type=ActionType.MOVE, direction=Direction.UP))
    return reward, done


def test_observation_includes_new_constraints() -> None:
    config = EnvironmentConfig(
        width=8,
        height=8,
        max_steps=80,
        max_orders=2,
        delivery_locations_per_order=2,
        priority_high_ratio=1.0,
        obstacle_density=0.2,
        dynamic_obstacles_enabled=True,
        dynamic_obstacle_ratio=0.5,
        traffic_density=0.1,
        battery_enabled=True,
        battery_capacity=20,
        battery_recharge_rate=5,
        seed=9,
    )
    env = LastMileDeliveryEnvironment(config=config)
    observation = env.reset()

    assert isinstance(observation.dynamic_obstacles, list)
    assert isinstance(observation.charging_stations, list)
    assert len(observation.pending_orders) > 0
    for order in observation.pending_orders:
        assert len(order.delivery_locations) >= 1
        assert order.priority in {OrderPriority.HIGH, OrderPriority.LOW}


def test_dynamic_obstacles_are_deterministic_on_reset() -> None:
    config = EnvironmentConfig(
        width=8,
        height=8,
        max_steps=50,
        max_orders=1,
        obstacle_density=0.2,
        dynamic_obstacles_enabled=True,
        dynamic_obstacle_ratio=1.0,
        seed=3,
    )
    env = LastMileDeliveryEnvironment(config=config)

    first = env.reset()
    second = env.reset()

    first_dynamic = sorted((p.x, p.y) for p in first.dynamic_obstacles)
    second_dynamic = sorted((p.x, p.y) for p in second.dynamic_obstacles)
    assert first_dynamic == second_dynamic


def test_battery_recharge_and_failure_flow() -> None:
    config = EnvironmentConfig(
        width=6,
        height=6,
        max_steps=30,
        max_orders=1,
        obstacle_density=0.0,
        traffic_density=0.0,
        battery_enabled=True,
        battery_capacity=3,
        battery_recharge_rate=3,
        seed=5,
    )
    env = LastMileDeliveryEnvironment(config=config)
    env.reset()

    _, _, done1, _ = env.step(SimulatorAction(action_type=ActionType.MOVE, direction=Direction.RIGHT))
    assert done1 is False

    _, _, done2, _ = env.step(SimulatorAction(action_type=ActionType.MOVE, direction=Direction.LEFT))
    assert done2 is False

    observation, reward_wait, done3, info_wait = env.step(SimulatorAction(action_type=ActionType.WAIT))
    assert done3 is False
    assert info_wait["message"] == "recharged"
    assert observation.battery_level == 3
    assert reward_wait > 0.0

    # Drain battery completely to trigger failure.
    env.step(SimulatorAction(action_type=ActionType.MOVE, direction=Direction.RIGHT))
    env.step(SimulatorAction(action_type=ActionType.MOVE, direction=Direction.RIGHT))
    _, reward_fail, done_fail, info_fail = env.step(SimulatorAction(action_type=ActionType.MOVE, direction=Direction.RIGHT))

    assert done_fail is True
    assert info_fail["battery_depleted"] is True
    assert reward_fail < 0.0


def test_high_priority_delivery_gets_bonus_reward() -> None:
    config = EnvironmentConfig(
        width=6,
        height=6,
        max_steps=120,
        max_orders=1,
        delivery_locations_per_order=2,
        priority_high_ratio=1.0,
        obstacle_density=0.0,
        traffic_density=0.0,
        battery_enabled=False,
        seed=11,
    )
    env = LastMileDeliveryEnvironment(config=config)
    observation = env.reset()

    order = observation.pending_orders[0]
    env.step(SimulatorAction(action_type=ActionType.ACCEPT_ORDER, order_id=order.order_id))

    pickup = order.pickup
    while True:
        current = env.current_observation().agent_location
        if current.x == pickup.x and current.y == pickup.y:
            break
        _, done = _move_towards(env, pickup.x, pickup.y)
        assert done is False

    # Move to the primary dropoff target.
    drop = order.delivery_locations[0]
    while True:
        current = env.current_observation().agent_location
        if current.x == drop.x and current.y == drop.y:
            break
        _, done = _move_towards(env, drop.x, drop.y)
        assert done is False

    _, reward, _, info = env.step(SimulatorAction(action_type=ActionType.DELIVER_ORDER))

    assert info["delivered_order_id"] == order.order_id
    assert reward >= 60.0