from collections.abc import Callable
from unittest.mock import ANY, MagicMock

import pytest
from bluesky.run_engine import RunEngine
from bluesky.simulators import RunEngineSimulator, assert_message_and_return_remaining
from bluesky.utils import Msg
from dodal.devices.aperturescatterguard import ApertureScatterguard, ApertureValue
from dodal.devices.motors import XYZStage
from dodal.devices.robot import SAMPLE_LOCATION_EMPTY, BartRobot
from dodal.devices.smargon import CombinedMove, Smargon, StubPosition
from ophyd_async.core import completed_status, get_mock_put, set_mock_value

from mx_bluesky.common.device_setup_plans.robot_load_unload import (
    prepare_for_robot_load,
    robot_unload,
)
from mx_bluesky.hyperion.external_interaction.callbacks.robot_actions.ispyb_callback import (
    RobotLoadISPyBCallback,
)

TEST_PUCK = 1
TEST_PIN = 2
TEST_SAMPLE_ID = 123456


@pytest.fixture
def loaded_robot(robot: BartRobot):
    set_mock_value(robot.current_pin, TEST_PIN)
    set_mock_value(robot.current_puck, TEST_PUCK)
    set_mock_value(robot.sample_id, TEST_SAMPLE_ID)
    return robot


# Remove when in bluesky proper, see https://github.com/bluesky/bluesky/issues/1924
def assert_messages_any_order(messages: list, predicates: list[Callable[[Msg], bool]]):
    remaining_predicates = set(predicates)

    for message in messages:
        matched = {p for p in remaining_predicates if p(message)}
        remaining_predicates -= matched

        if not remaining_predicates:
            break

    assert not remaining_predicates, f"Unmatched predicates: {remaining_predicates}"
    return messages


async def test_when_prepare_for_robot_load_called_then_moves_as_expected(
    aperture_scatterguard: ApertureScatterguard,
    smargon: Smargon,
    run_engine: RunEngine,
):
    smargon.stub_offsets.set = MagicMock(side_effect=lambda _: completed_status())
    get_mock_put(aperture_scatterguard.selected_aperture).reset_mock()

    set_mock_value(smargon.x.user_setpoint, 10)
    set_mock_value(smargon.z.user_setpoint, 5)
    set_mock_value(smargon.omega.user_setpoint, 90)

    run_engine(prepare_for_robot_load(aperture_scatterguard, smargon))

    assert await smargon.x.user_setpoint.get_value() == 0
    assert await smargon.z.user_setpoint.get_value() == 0
    assert await smargon.omega.user_setpoint.get_value() == 0

    smargon.stub_offsets.set.assert_called_once_with(StubPosition.RESET_TO_ROBOT_LOAD)  # type: ignore
    get_mock_put(aperture_scatterguard.selected_aperture).assert_called_once_with(
        ApertureValue.OUT_OF_BEAM
    )


async def test_when_robot_unload_called_then_sample_area_prepared_before_load(
    robot: BartRobot,
    smargon: Smargon,
    aperture_scatterguard: ApertureScatterguard,
    lower_gonio: XYZStage,
    sim_run_engine: RunEngineSimulator,
):
    msgs = sim_run_engine.simulate_plan(
        robot_unload(robot, smargon, aperture_scatterguard, lower_gonio, "")
    )

    msgs = assert_message_and_return_remaining(
        msgs,
        lambda msg: msg.command == "set"
        and msg.obj.name == aperture_scatterguard.selected_aperture.name
        and msg.args[0] == ApertureValue.OUT_OF_BEAM,
    )

    assert_message_and_return_remaining(
        msgs,
        lambda msg: msg.command == "set"
        and msg.obj.name == "gonio"
        and msg.args[0] == CombinedMove(x=0, y=0, z=0, omega=0, chi=0, phi=0),
    )

    assert_message_and_return_remaining(
        msgs,
        lambda msg: msg.command == "set"
        and msg.obj is robot
        and msg.args[0] == SAMPLE_LOCATION_EMPTY,
    )
    assert_message_and_return_remaining(
        msgs,
        lambda msg: msg.command == "wait"
        and msg.kwargs.get("group") == msgs[0].kwargs["group"],
    )


async def test_given_lower_gonio_needs_moving_then_it_is_homed_before_unload_and_put_back_after(
    robot: BartRobot,
    smargon: Smargon,
    aperture_scatterguard: ApertureScatterguard,
    lower_gonio: XYZStage,
    sim_run_engine: RunEngineSimulator,
):
    # Replace when https://github.com/bluesky/bluesky/issues/1906 is fixed
    def locate_gonio(_):
        return {"readback": 0.1}

    sim_run_engine.add_handler("locate", locate_gonio, lower_gonio.x.name)
    sim_run_engine.add_handler("locate", locate_gonio, lower_gonio.y.name)
    sim_run_engine.add_handler("locate", locate_gonio, lower_gonio.z.name)

    sim_run_engine.add_read_handler_for(robot.sample_id, 1000)

    msgs = sim_run_engine.simulate_plan(
        robot_unload(robot, smargon, aperture_scatterguard, lower_gonio, "")
    )

    msgs = assert_messages_any_order(
        msgs,
        [
            lambda msg: msg.command == "set"
            and msg.obj.name == f"lower_gonio-{axis}"
            and msg.args[0] == 0
            for axis in ["x", "y", "z"]
        ],
    )

    msgs = assert_message_and_return_remaining(
        msgs,
        lambda msg: msg.command == "set"
        and msg.obj is robot
        and msg.args[0] == SAMPLE_LOCATION_EMPTY,
    )

    msgs = assert_messages_any_order(
        msgs,
        [
            lambda msg: msg.command == "set"
            and msg.obj.name == f"lower_gonio-{axis}"
            and msg.args[0] == 0.1
            for axis in ["x", "y", "z"]
        ],
    )


def test_when_unload_plan_run_then_initial_unload_ispyb_deposition_made(
    run_engine: RunEngine,
    robot: BartRobot,
    smargon: Smargon,
    aperture_scatterguard: ApertureScatterguard,
    lower_gonio: XYZStage,
):
    callback = RobotLoadISPyBCallback()
    callback.expeye = (mock_expeye := MagicMock())
    run_engine.subscribe(callback)

    set_mock_value(robot.sample_id, expected_sample_id := 1234)

    run_engine(
        robot_unload(robot, smargon, aperture_scatterguard, lower_gonio, "cm37235-2")
    )

    mock_expeye.start_robot_action.assert_called_once_with(
        "UNLOAD", "cm37235", 2, expected_sample_id
    )


def test_when_unload_plan_run_then_full_ispyb_deposition_made(
    run_engine: RunEngine,
    robot: BartRobot,
    smargon: Smargon,
    aperture_scatterguard: ApertureScatterguard,
    lower_gonio: XYZStage,
):
    callback = RobotLoadISPyBCallback()
    callback.expeye = (mock_expeye := MagicMock())
    run_engine.subscribe(callback)

    set_mock_value(robot.sample_id, expected_sample_id := 1234)
    set_mock_value(robot.current_pin, expected_pin := 12)
    set_mock_value(robot.current_puck, expected_puck := 45)
    set_mock_value(robot.barcode, expected_barcode := "BARCODE")

    action_id = 1098
    mock_expeye.start_robot_action.return_value = action_id

    run_engine(
        robot_unload(robot, smargon, aperture_scatterguard, lower_gonio, "cm37235-2")
    )

    mock_expeye.start_robot_action.assert_called_once_with(
        "UNLOAD", "cm37235", 2, expected_sample_id
    )
    mock_expeye.update_robot_action.assert_called_once_with(
        action_id,
        {
            "sampleBarcode": expected_barcode,
            "containerLocation": expected_pin,
            "dewarLocation": expected_puck,
        },
    )
    mock_expeye.end_robot_action.assert_called_once_with(action_id, "success", "OK")


def test_when_unload_plan_fails_then_error_deposited_in_ispyb(
    run_engine: RunEngine,
    loaded_robot: BartRobot,
    smargon: Smargon,
    aperture_scatterguard: ApertureScatterguard,
    lower_gonio: XYZStage,
):
    class TestError(Exception): ...

    callback = RobotLoadISPyBCallback()
    callback.expeye = (mock_expeye := MagicMock())
    run_engine.subscribe(callback)
    loaded_robot.set = MagicMock(side_effect=TestError("Bad Error"))

    action_id = 1098
    mock_expeye.start_robot_action.return_value = action_id

    with pytest.raises(TestError):
        run_engine(
            robot_unload(
                loaded_robot, smargon, aperture_scatterguard, lower_gonio, "cm37235-2"
            )
        )

    mock_expeye.start_robot_action.assert_called_once_with("UNLOAD", "cm37235", 2, ANY)
    mock_expeye.end_robot_action.assert_called_once_with(action_id, "fail", "Bad Error")
