"""MovementHandler command subject."""

from typing import NamedTuple, Optional

import pytest
from decoy import Decoy

from opentrons_shared_data.errors.exceptions import PositionUnknownError

from opentrons.hardware_control import API as HardwareAPI
from opentrons.hardware_control.types import CriticalPoint
from opentrons.motion_planning import Waypoint
from opentrons.protocol_engine.execution.equipment import EquipmentHandler
from opentrons.protocol_engine.execution.gantry_mover import GantryMover
from opentrons.protocol_engine.execution.heater_shaker_movement_flagger import (
    HeaterShakerMovementFlagger,
)
from opentrons.protocol_engine.execution.movement import MovementHandler
from opentrons.protocol_engine.execution.thermocycler_movement_flagger import (
    ThermocyclerMovementFlagger,
)
from opentrons.protocol_engine.state.motion import PipetteLocationData
from opentrons.protocol_engine.state.state import (
    StateStore,
)
from opentrons.protocol_engine.types import (
    AddressableOffsetVector,
    CurrentWell,
    DeckPoint,
    DeckSlotLocation,
    MotorAxis,
    MovementAxis,
    WellLocation,
    WellOffset,
    WellOrigin,
)
from opentrons.types import DeckSlotName, Mount, MountType, Point


@pytest.fixture
def hardware_api(decoy: Decoy) -> HardwareAPI:
    """Get a mock in the shape of a HardwareAPI."""
    return decoy.mock(cls=HardwareAPI)


@pytest.fixture
def state_store(decoy: Decoy) -> StateStore:
    """Get a mock in the shape of a StateStore."""
    return decoy.mock(cls=StateStore)


@pytest.fixture
def thermocycler_movement_flagger(decoy: Decoy) -> ThermocyclerMovementFlagger:
    """Get a mock in the shape of a ThermocyclerMovementFlagger."""
    return decoy.mock(cls=ThermocyclerMovementFlagger)


@pytest.fixture
def heater_shaker_movement_flagger(decoy: Decoy) -> HeaterShakerMovementFlagger:
    """Get a mock in the shape of a HeaterShakerMovementFlagger."""
    return decoy.mock(cls=HeaterShakerMovementFlagger)


@pytest.fixture
def mock_gantry_mover(decoy: Decoy) -> GantryMover:
    """Get a mock in the shape of a GantryMover."""
    return decoy.mock(cls=GantryMover)


@pytest.fixture
def mock_equipment_handler(decoy: Decoy) -> EquipmentHandler:
    """Get a mock in the shape of an EquipmentHandler."""
    return decoy.mock(cls=EquipmentHandler)


@pytest.fixture
def subject(
    state_store: StateStore,
    hardware_api: HardwareAPI,
    thermocycler_movement_flagger: ThermocyclerMovementFlagger,
    heater_shaker_movement_flagger: HeaterShakerMovementFlagger,
    mock_gantry_mover: GantryMover,
    mock_equipment_handler: EquipmentHandler,
) -> MovementHandler:
    """Create a MovementHandler with its dependencies mocked out."""
    return MovementHandler(
        state_store=state_store,
        hardware_api=hardware_api,
        thermocycler_movement_flagger=thermocycler_movement_flagger,
        heater_shaker_movement_flagger=heater_shaker_movement_flagger,
        gantry_mover=mock_gantry_mover,
        equipment=mock_equipment_handler,
    )


async def test_move_to_well(
    decoy: Decoy,
    state_store: StateStore,
    thermocycler_movement_flagger: ThermocyclerMovementFlagger,
    heater_shaker_movement_flagger: HeaterShakerMovementFlagger,
    mock_gantry_mover: GantryMover,
    subject: MovementHandler,
) -> None:
    """Move requests should call hardware controller with movement data."""
    well_location = WellLocation(
        origin=WellOrigin.BOTTOM,
        offset=WellOffset(x=0, y=0, z=1),
    )
    decoy.when(state_store.labware.get_location(labware_id="labware-id")).then_return(
        DeckSlotLocation(slotName=DeckSlotName.SLOT_1)
    )
    decoy.when(
        state_store.modules.get_heater_shaker_movement_restrictors()
    ).then_return([])

    decoy.when(state_store.geometry.get_ancestor_slot_name("labware-id")).then_return(
        DeckSlotName.SLOT_1
    )

    decoy.when(state_store.pipettes.get_channels("pipette-id")).then_return(1)
    decoy.when(state_store.labware.is_tiprack("labware-id")).then_return(False)

    decoy.when(
        state_store.motion.get_pipette_location(
            pipette_id="pipette-id",
            current_location=None,
        )
    ).then_return(
        PipetteLocationData(
            mount=MountType.LEFT,
            critical_point=CriticalPoint.FRONT_NOZZLE,
        )
    )

    decoy.when(
        await mock_gantry_mover.get_position(
            pipette_id="pipette-id",
        )
    ).then_return(Point(1, 1, 1))

    decoy.when(mock_gantry_mover.get_max_travel_z(pipette_id="pipette-id")).then_return(
        42.0
    )

    decoy.when(
        state_store.pipettes.get_movement_speed(
            pipette_id="pipette-id", requested_speed=45.6
        )
    ).then_return(39339.5)

    decoy.when(
        state_store.motion.get_movement_waypoints_to_well(
            origin=Point(1, 1, 1),
            origin_cp=CriticalPoint.FRONT_NOZZLE,
            max_travel_z=42.0,
            pipette_id="pipette-id",
            labware_id="labware-id",
            well_name="B2",
            well_location=well_location,
            current_well=None,
            force_direct=True,
            minimum_z_height=12.3,
            operation_volume=None,
            offset_pipette_for_reservoir_subwells=False,
        )
    ).then_return(
        [Waypoint(Point(1, 2, 3), CriticalPoint.XY_CENTER), Waypoint(Point(4, 5, 6))]
    )

    decoy.when(
        await mock_gantry_mover.move_to(
            pipette_id="pipette-id",
            waypoints=[
                Waypoint(Point(1, 2, 3), CriticalPoint.XY_CENTER),
                Waypoint(Point(4, 5, 6)),
            ],
            speed=39339.5,
        ),
    ).then_return(Point(4, 5, 6))

    result = await subject.move_to_well(
        pipette_id="pipette-id",
        labware_id="labware-id",
        well_name="B2",
        well_location=well_location,
        force_direct=True,
        minimum_z_height=12.3,
        speed=45.6,
    )

    assert result == Point(x=4, y=5, z=6)

    decoy.verify(
        await thermocycler_movement_flagger.ensure_labware_in_open_thermocycler(
            labware_parent=DeckSlotLocation(slotName=DeckSlotName.SLOT_1)
        ),
        heater_shaker_movement_flagger.raise_if_movement_restricted(
            hs_movement_restrictors=[],
            destination_slot=1,
            is_multi_channel=False,
            destination_is_tip_rack=False,
        ),
    )


async def test_move_to_well_from_starting_location(
    decoy: Decoy,
    state_store: StateStore,
    thermocycler_movement_flagger: ThermocyclerMovementFlagger,
    heater_shaker_movement_flagger: HeaterShakerMovementFlagger,
    mock_gantry_mover: GantryMover,
    subject: MovementHandler,
) -> None:
    """It should be able to move to a well from a start location."""
    well_location = WellLocation(
        origin=WellOrigin.BOTTOM,
        offset=WellOffset(x=0, y=0, z=1),
    )

    current_well = CurrentWell(
        pipette_id="pipette-id",
        labware_id="labware-id",
        well_name="B2",
    )
    decoy.when(state_store.labware.get_location(labware_id="labware-id")).then_return(
        DeckSlotLocation(slotName=DeckSlotName.SLOT_1)
    )
    decoy.when(
        state_store.modules.get_heater_shaker_movement_restrictors()
    ).then_return([])

    decoy.when(state_store.geometry.get_ancestor_slot_name("labware-id")).then_return(
        DeckSlotName.SLOT_1
    )

    decoy.when(state_store.pipettes.get_channels("pipette-id")).then_return(1)
    decoy.when(state_store.labware.is_tiprack("labware-id")).then_return(False)

    decoy.when(
        state_store.motion.get_pipette_location(
            pipette_id="pipette-id",
            current_location=current_well,
        )
    ).then_return(
        PipetteLocationData(
            mount=MountType.RIGHT,
            critical_point=CriticalPoint.XY_CENTER,
        )
    )

    decoy.when(
        await mock_gantry_mover.get_position(
            pipette_id="pipette-id",
        )
    ).then_return(Point(1, 2, 5))

    decoy.when(mock_gantry_mover.get_max_travel_z(pipette_id="pipette-id")).then_return(
        42.0
    )

    decoy.when(
        state_store.motion.get_movement_waypoints_to_well(
            current_well=current_well,
            origin=Point(1, 2, 5),
            origin_cp=CriticalPoint.XY_CENTER,
            max_travel_z=42.0,
            pipette_id="pipette-id",
            labware_id="labware-id",
            well_name="B2",
            well_location=well_location,
            force_direct=False,
            minimum_z_height=None,
            operation_volume=None,
            offset_pipette_for_reservoir_subwells=False,
        )
    ).then_return([Waypoint(Point(1, 2, 3), CriticalPoint.XY_CENTER)])

    decoy.when(
        state_store.pipettes.get_movement_speed(
            pipette_id="pipette-id", requested_speed=None
        )
    ).then_return(39339.5)

    decoy.when(
        await mock_gantry_mover.move_to(
            pipette_id="pipette-id",
            waypoints=[Waypoint(Point(1, 2, 3), CriticalPoint.XY_CENTER)],
            speed=39339.5,
        ),
    ).then_return(Point(4, 5, 6))

    result = await subject.move_to_well(
        pipette_id="pipette-id",
        labware_id="labware-id",
        well_name="B2",
        well_location=well_location,
        current_well=current_well,
    )

    assert result == Point(4, 5, 6)

    decoy.verify(
        await thermocycler_movement_flagger.ensure_labware_in_open_thermocycler(
            labware_parent=DeckSlotLocation(slotName=DeckSlotName.SLOT_1)
        ),
        heater_shaker_movement_flagger.raise_if_movement_restricted(
            hs_movement_restrictors=[],
            destination_slot=1,
            is_multi_channel=False,
            destination_is_tip_rack=False,
        ),
    )


@pytest.mark.parametrize(
    "stay_at_max_z,z_extra_offset",
    [(True, None), (True, 5.0), (False, None), (False, 5.0)],
)
async def test_move_to_addressable_area(
    decoy: Decoy,
    state_store: StateStore,
    thermocycler_movement_flagger: ThermocyclerMovementFlagger,
    heater_shaker_movement_flagger: HeaterShakerMovementFlagger,
    mock_gantry_mover: GantryMover,
    subject: MovementHandler,
    stay_at_max_z: bool,
    z_extra_offset: Optional[float],
) -> None:
    """Move requests should call hardware controller with movement data."""
    decoy.when(
        state_store.modules.get_heater_shaker_movement_restrictors()
    ).then_return([])

    decoy.when(
        state_store.addressable_areas.get_addressable_area_base_slot("area-name")
    ).then_return(DeckSlotName.SLOT_1)

    decoy.when(state_store.pipettes.get_channels("pipette-id")).then_return(1)

    decoy.when(
        state_store.motion.get_pipette_location(
            pipette_id="pipette-id",
            current_location=None,
        )
    ).then_return(
        PipetteLocationData(
            mount=MountType.LEFT,
            critical_point=CriticalPoint.FRONT_NOZZLE,
        )
    )

    decoy.when(
        await mock_gantry_mover.get_position(
            pipette_id="pipette-id",
        )
    ).then_return(Point(1, 1, 1))

    decoy.when(mock_gantry_mover.get_max_travel_z(pipette_id="pipette-id")).then_return(
        42.0
    )

    decoy.when(
        state_store.pipettes.get_movement_speed(
            pipette_id="pipette-id", requested_speed=45.6
        )
    ).then_return(39339.5)

    decoy.when(
        state_store.motion.get_movement_waypoints_to_addressable_area(
            addressable_area_name="area-name",
            offset=AddressableOffsetVector(x=9, y=8, z=7),
            origin=Point(1, 1, 1),
            origin_cp=CriticalPoint.FRONT_NOZZLE,
            max_travel_z=42.0,
            force_direct=True,
            minimum_z_height=12.3,
            stay_at_max_travel_z=stay_at_max_z,
            ignore_tip_configuration=False,
            max_travel_z_extra_margin=z_extra_offset,
        )
    ).then_return(
        [Waypoint(Point(1, 2, 3), CriticalPoint.XY_CENTER), Waypoint(Point(4, 5, 6))]
    )

    decoy.when(
        await mock_gantry_mover.move_to(
            pipette_id="pipette-id",
            waypoints=[
                Waypoint(Point(1, 2, 3), CriticalPoint.XY_CENTER),
                Waypoint(Point(4, 5, 6)),
            ],
            speed=39339.5,
        ),
    ).then_return(Point(4, 5, 6))

    result = await subject.move_to_addressable_area(
        pipette_id="pipette-id",
        addressable_area_name="area-name",
        offset=AddressableOffsetVector(x=9, y=8, z=7),
        force_direct=True,
        minimum_z_height=12.3,
        speed=45.6,
        stay_at_highest_possible_z=stay_at_max_z,
        ignore_tip_configuration=False,
        highest_possible_z_extra_offset=z_extra_offset,
    )

    assert result == Point(x=4, y=5, z=6)

    decoy.verify(
        heater_shaker_movement_flagger.raise_if_movement_restricted(
            hs_movement_restrictors=[],
            destination_slot=1,
            is_multi_channel=False,
            destination_is_tip_rack=False,
        )
    )


class MoveRelativeSpec(NamedTuple):
    """Test data for move_relative."""

    axis: MovementAxis
    expected_delta: Point
    distance: float = 42.0


@pytest.mark.parametrize(
    MoveRelativeSpec._fields,
    [
        MoveRelativeSpec(
            axis=MovementAxis.X,
            expected_delta=Point(x=42.0, y=0, z=0),
        ),
        MoveRelativeSpec(
            axis=MovementAxis.Y,
            expected_delta=Point(x=0, y=42.0, z=0),
        ),
        MoveRelativeSpec(
            axis=MovementAxis.Z,
            expected_delta=Point(x=0, y=0, z=42.0),
        ),
    ],
)
async def test_move_relative(
    decoy: Decoy,
    state_store: StateStore,
    mock_gantry_mover: GantryMover,
    subject: MovementHandler,
    axis: MovementAxis,
    expected_delta: Point,
    distance: float,
) -> None:
    """Test that move_relative triggers a relative move with the HardwareAPI."""
    decoy.when(
        await mock_gantry_mover.move_relative(
            pipette_id="pipette-id",
            delta=expected_delta,
            speed=39339.5,
        )
    ).then_return(Point(x=1, y=2, z=3))

    decoy.when(
        state_store.pipettes.get_movement_speed(pipette_id="pipette-id")
    ).then_return(39339.5)

    result = await subject.move_relative(
        pipette_id="pipette-id",
        axis=axis,
        distance=distance,
    )

    assert result == Point(x=1, y=2, z=3)


async def test_move_to_coordinates(
    decoy: Decoy,
    state_store: StateStore,
    mock_gantry_mover: GantryMover,
    subject: MovementHandler,
) -> None:
    """Test that move_to_coordinates correctly calls api.move_to."""
    current_position = Point(4.44, 5.55, 6.66)
    destination_deck = DeckPoint(x=1.11, y=2.22, z=3.33)
    destination_point = Point(1.11, 2.22, 3.33)

    planned_waypoint_1 = Waypoint(position=Point(3, 1, 4), critical_point=None)
    planned_waypoint_2 = Waypoint(
        position=Point(1, 5, 9), critical_point=CriticalPoint.XY_CENTER
    )

    decoy.when(
        state_store.motion.get_pipette_location(
            pipette_id="pipette-id",
        )
    ).then_return(
        PipetteLocationData(
            mount=MountType.RIGHT,
            critical_point=CriticalPoint.XY_CENTER,
        )
    )

    decoy.when(
        await mock_gantry_mover.get_position(pipette_id="pipette-id")
    ).then_return(current_position)

    decoy.when(mock_gantry_mover.get_max_travel_z(pipette_id="pipette-id")).then_return(
        5678
    )

    decoy.when(
        state_store.motion.get_movement_waypoints_to_coords(
            origin=current_position,
            dest=destination_point,
            max_travel_z=5678,
            direct=True,
            additional_min_travel_z=1234,
        )
    ).then_return([planned_waypoint_1, planned_waypoint_2])

    decoy.when(
        state_store.pipettes.get_movement_speed(
            pipette_id="pipette-id", requested_speed=567
        )
    ).then_return(39339.5)

    decoy.when(
        await mock_gantry_mover.move_to(
            pipette_id="pipette-id",
            waypoints=[planned_waypoint_1, planned_waypoint_2],
            speed=39339.5,
        )
    ).then_return(Point(x=1, y=5, z=9))

    result = await subject.move_to_coordinates(
        pipette_id="pipette-id",
        deck_coordinates=destination_deck,
        direct=True,
        additional_min_travel_z=1234,
        speed=567,
    )

    assert result == Point(x=1, y=5, z=9)


async def test_retract_axis(
    decoy: Decoy,
    state_store: StateStore,
    mock_gantry_mover: GantryMover,
    subject: MovementHandler,
) -> None:
    """It should delegate to gantry mover to retract the specified axis."""
    await subject.retract_axis(axis=MotorAxis.RIGHT_Z)

    decoy.verify(await mock_gantry_mover.retract_axis(MotorAxis.RIGHT_Z), times=1)


async def test_check_valid_position(
    decoy: Decoy,
    subject: MovementHandler,
    hardware_api: HardwareAPI,
) -> None:
    """It should check for an exception to determine if the position is ok."""
    decoy.when(
        await hardware_api.gantry_position(mount=Mount.LEFT, fail_on_not_homed=True)
    ).then_return(Point(0, 0, 0))
    decoy.when(
        await hardware_api.gantry_position(mount=Mount.RIGHT, fail_on_not_homed=True)
    ).then_raise(PositionUnknownError())

    assert await subject.check_for_valid_position(MountType.LEFT)
    assert not await subject.check_for_valid_position(MountType.RIGHT)
