from functools import partial
from unittest.mock import MagicMock, call, patch

import numpy as np
import pytest
from bluesky import Msg
from bluesky.run_engine import RunEngine
from bluesky.simulators import RunEngineSimulator, assert_message_and_return_remaining
from bluesky.utils import MsgGenerator
from dodal.devices.aperturescatterguard import ApertureScatterguard, ApertureValue
from dodal.devices.attenuator.attenuator import BinaryFilterAttenuator
from dodal.devices.backlight import Backlight
from dodal.devices.beamlines.i04.beamsize import Beamsize
from dodal.devices.beamlines.i04.transfocator import Transfocator
from dodal.devices.common_dcm import DoubleCrystalMonochromator
from dodal.devices.detector.detector_motion import DetectorMotion
from dodal.devices.eiger import EigerDetector
from dodal.devices.fast_grid_scan import (
    ZebraFastGridScanThreeD,
)
from dodal.devices.flux import Flux
from dodal.devices.mx_phase1.beamstop import Beamstop
from dodal.devices.oav.oav_detector import OAV
from dodal.devices.oav.pin_image_recognition import PinTipDetection
from dodal.devices.robot import BartRobot
from dodal.devices.s4_slit_gaps import S4SlitGaps
from dodal.devices.smargon import Smargon
from dodal.devices.synchrotron import Synchrotron
from dodal.devices.undulator import UndulatorInKeV
from dodal.devices.xbpm_feedback import XBPMFeedback
from dodal.devices.zebra.zebra import Zebra
from dodal.devices.zebra.zebra_controlled_shutter import ZebraShutter
from dodal.devices.zocalo import ZocaloResults
from dodal.devices.zocalo.zocalo_results import (
    ZOCALO_STAGE_GROUP,
)
from ophyd_async.core import get_mock_put, set_mock_value

from mx_bluesky.beamlines.i04.experiment_plans.i04_grid_detect_then_xray_centre_plan import (
    DEFAULT_XRC_BEAMSIZE_MICRONS,
    I04AutoXrcParams,
    _get_grid_common_params,
    get_ready_for_oav_and_close_shutter,
    i04_default_grid_detect_and_xray_centre,
)
from mx_bluesky.common.parameters.constants import PlanNameConstants
from mx_bluesky.common.parameters.gridscan import (
    GridCommon,
    SpecifiedThreeDGridScan,
)
from mx_bluesky.common.utils.exceptions import CrystalNotFoundError
from mx_bluesky.common.utils.xrc_result import XRayCentreResult
from tests.conftest import TEST_RESULT_LARGE, fake_generator, simulate_xrc_result
from tests.unit_tests.common.experiment_plans.test_common_flyscan_xray_centre_plan import (
    CompleteError,
)


class CustomError(Exception): ...


EXPECTED_WAVELENGTH = 0.95373


@pytest.fixture
def i04_grid_detect_then_xrc_default_params(
    aperture_scatterguard: ApertureScatterguard,
    attenuator: BinaryFilterAttenuator,
    backlight: Backlight,
    beamsize: Beamsize,
    beamstop_phase1: Beamstop,
    dcm: DoubleCrystalMonochromator,
    zebra_fast_grid_scan: ZebraFastGridScanThreeD,
    flux: Flux,
    oav: OAV,
    pin_tip_detection_with_found_pin: PinTipDetection,
    s4_slit_gaps: S4SlitGaps,
    undulator: UndulatorInKeV,
    xbpm_feedback: XBPMFeedback,
    zebra: Zebra,
    robot: BartRobot,
    sample_shutter: ZebraShutter,
    eiger: EigerDetector,
    synchrotron: Synchrotron,
    zocalo: ZocaloResults,
    smargon: Smargon,
    detector_motion: DetectorMotion,
    transfocator: Transfocator,
    tmp_path,
):
    entry_params = I04AutoXrcParams(
        sample_id=1,
        file_name="filename",
        visit="cm40607-5",
        detector_distance_mm=264.5,
        storage_directory=str(tmp_path),
    )

    set_mock_value(dcm.wavelength_in_a.user_readback, EXPECTED_WAVELENGTH)
    return partial(
        i04_default_grid_detect_and_xray_centre,
        parameters=entry_params,
        aperture_scatterguard=aperture_scatterguard,
        attenuator=attenuator,
        backlight=backlight,
        beamstop=beamstop_phase1,
        beamsize=beamsize,
        dcm=dcm,
        zebra_fast_grid_scan=zebra_fast_grid_scan,
        flux=flux,
        oav=oav,
        pin_tip_detection=pin_tip_detection_with_found_pin,
        s4_slit_gaps=s4_slit_gaps,
        undulator=undulator,
        xbpm_feedback=xbpm_feedback,
        zebra=zebra,
        robot=robot,
        sample_shutter=sample_shutter,
        eiger=eiger,
        synchrotron=synchrotron,
        zocalo=zocalo,
        smargon=smargon,
        detector_motion=detector_motion,
        transfocator=transfocator,
    )


@patch(
    "mx_bluesky.beamlines.i04.experiment_plans.i04_grid_detect_then_xray_centre_plan.setup_beamline_for_oav",
    autospec=True,
)
def test_get_ready_for_oav_and_close_shutter_closes_shutter_and_calls_setup_for_oav_plan(
    mock_setup_beamline_for_oav: MagicMock,
    sim_run_engine: RunEngineSimulator,
    grid_detect_xrc_devices,
):
    mock_setup_beamline_for_oav.return_value = iter([Msg("setup_beamline_for_oav")])

    msgs = sim_run_engine.simulate_plan(
        get_ready_for_oav_and_close_shutter(
            grid_detect_xrc_devices.gonio,
            grid_detect_xrc_devices.backlight,
            grid_detect_xrc_devices.aperture_scatterguard,
            grid_detect_xrc_devices.detector_motion,
        )
    )
    msgs = assert_message_and_return_remaining(
        msgs, predicate=lambda msg: msg.command == "wait"
    )

    msgs = assert_message_and_return_remaining(
        msgs, predicate=lambda msg: msg.command == "setup_beamline_for_oav"
    )
    msgs = assert_message_and_return_remaining(
        msgs,
        predicate=lambda msg: msg.command == "set"
        and msg.obj.name == "detector_motion-shutter"
        and msg.args[0] == 0,
    )
    msgs = assert_message_and_return_remaining(
        msgs, predicate=lambda msg: msg.command == "wait"
    )


@pytest.mark.parametrize(
    "udc",
    [(True), (False)],
)
@patch(
    "mx_bluesky.beamlines.i04.experiment_plans.i04_grid_detect_then_xray_centre_plan.get_ready_for_oav_and_close_shutter",
    autospec=True,
)
@patch(
    "mx_bluesky.beamlines.i04.experiment_plans.i04_grid_detect_then_xray_centre_plan.grid_detect_then_xray_centre",
    autospec=True,
)
@patch(
    "mx_bluesky.beamlines.i04.experiment_plans.i04_grid_detect_then_xray_centre_plan.setup_beamline_for_oav",
    autospec=True,
)
@patch(
    "mx_bluesky.beamlines.i04.experiment_plans.i04_grid_detect_then_xray_centre_plan.create_gridscan_callbacks",
    autospec=True,
)
@patch(
    "mx_bluesky.beamlines.i04.experiment_plans.i04_grid_detect_then_xray_centre_plan.get_results_and_move_to_xtal",
    new=MagicMock(),
)
def test_i04_grid_detect_then_xrc_closes_shutter_and_tidies_if_not_udc(
    mock_create_gridscan_callbacks: MagicMock,
    mock_setup_beamline_for_oav: MagicMock,
    mock_grid_detect_then_xray_centre: MagicMock,
    mock_get_ready_for_oav_and_close_shutter: MagicMock,
    run_engine: RunEngine,
    i04_grid_detect_then_xrc_default_params: partial[MsgGenerator],
    udc: bool,
    give_grid_common_specific_grid,
):
    run_engine(
        i04_grid_detect_then_xrc_default_params(
            udc=udc,
        )
    )

    call_count = 0 if udc else 1

    assert mock_get_ready_for_oav_and_close_shutter.call_count == call_count


@patch("bluesky.plan_stubs.sleep", autospec=True)
@patch(
    "mx_bluesky.common.experiment_plans.inner_plans.do_fgs.check_topup_and_wait_if_necessary",
)
@patch(
    "mx_bluesky.common.experiment_plans.common_grid_detect_then_xray_centre_plan.grid_detection_plan",
)
@patch(
    "mx_bluesky.common.experiment_plans.common_grid_detect_then_xray_centre_plan.GridDetectionCallback",
)
@patch(
    "mx_bluesky.common.experiment_plans.common_grid_detect_then_xray_centre_plan.create_parameters_for_flyscan_xray_centre",
)
@patch(
    "mx_bluesky.beamlines.i04.experiment_plans.i04_grid_detect_then_xray_centre_plan.fix_transmission_and_exposure_time_for_current_wavelength"
)
@patch(
    "mx_bluesky.common.preprocessors.preprocessors.check_and_pause_feedback",
    autospec=True,
)
@patch(
    "mx_bluesky.beamlines.i04.experiment_plans.i04_grid_detect_then_xray_centre_plan.get_results_and_move_to_xtal",
    new=MagicMock(),
)
def test_i04_default_grid_detect_and_xray_centre_sets_transmission_and_triggers_xbpm_feedback_before_run(
    mock_pause_feedback: MagicMock,
    mock_fix_transmission_and_exp_time: MagicMock,
    mock_create_parameters: MagicMock,
    mock_grid_detection_callback: MagicMock,
    mock_grid_detection_plan: MagicMock,
    mock_check_topup: MagicMock,
    mock_wait: MagicMock,
    sim_run_engine: RunEngineSimulator,
    zocalo: ZocaloResults,
    hyperion_fgs_params,
    i04_grid_detect_then_xrc_default_params: partial[MsgGenerator],
):
    desired_transmission = 0.4
    mock_fix_transmission_and_exp_time.return_value = (desired_transmission, 1)
    mock_create_parameters.return_value = hyperion_fgs_params
    simulate_xrc_result(
        sim_run_engine,
        zocalo,
        TEST_RESULT_LARGE,
    )

    msgs = sim_run_engine.simulate_plan(
        i04_grid_detect_then_xrc_default_params(),
    )

    msgs = assert_message_and_return_remaining(
        msgs,
        lambda msg: msg.command == "set"
        and msg.obj.name == "attenuator"
        and msg.args == (desired_transmission,),
    )

    msgs = assert_message_and_return_remaining(
        msgs,
        lambda msg: msg.command == "trigger" and msg.obj.name == "xbpm_feedback",
    )
    msgs = assert_message_and_return_remaining(
        msgs,
        lambda msg: msg.command == "open_run"
        and msg.run == PlanNameConstants.GRIDSCAN_OUTER,
    )

    msgs = assert_message_and_return_remaining(
        msgs,
        lambda msg: msg.command == "close_run"
        and msg.run == PlanNameConstants.GRIDSCAN_OUTER,
    )

    mock_pause_feedback.assert_not_called()


@patch(
    "mx_bluesky.beamlines.i04.experiment_plans.i04_grid_detect_then_xray_centre_plan.create_gridscan_callbacks",
    autospec=True,
)
@patch(
    "mx_bluesky.common.experiment_plans.common_grid_detect_then_xray_centre_plan.grid_detection_plan",
)
@patch(
    "mx_bluesky.common.experiment_plans.common_grid_detect_then_xray_centre_plan.move_aperture_if_required",
)
@patch(
    "mx_bluesky.common.experiment_plans.common_grid_detect_then_xray_centre_plan.GridDetectionCallback",
)
@patch(
    "mx_bluesky.common.experiment_plans.common_grid_detect_then_xray_centre_plan.create_parameters_for_flyscan_xray_centre",
)
@patch(
    "mx_bluesky.common.experiment_plans.common_flyscan_xray_centre_plan.run_gridscan",
)
@patch(
    "dodal.plans.preprocessors.verify_undulator_gap.verify_undulator_gap",
)
def test_i04_default_grid_detect_and_xray_centre_does_undulator_check_before_collection(
    mock_verify_gap: MagicMock,
    mock_run_gridscan: MagicMock,
    mock_create_parameters: MagicMock,
    mock_grid_params_callback: MagicMock,
    mock_move_aperture_if_required: MagicMock,
    mock_grid_detection_plan: MagicMock,
    mock_create_gridscan_callbacks: MagicMock,
    run_engine: RunEngine,
    hyperion_fgs_params,
    i04_grid_detect_then_xrc_default_params: partial[MsgGenerator],
):
    mock_create_parameters.return_value = hyperion_fgs_params
    mock_run_gridscan.side_effect = CompleteError
    with pytest.raises(CompleteError):
        run_engine(i04_grid_detect_then_xrc_default_params())

    mock_verify_gap.assert_called_once()


@patch(
    "mx_bluesky.beamlines.i04.experiment_plans.i04_grid_detect_then_xray_centre_plan.get_ready_for_oav_and_close_shutter",
    autospec=True,
)
@patch(
    "mx_bluesky.beamlines.i04.experiment_plans.i04_grid_detect_then_xray_centre_plan.create_gridscan_callbacks",
    autospec=True,
)
def test_i04_grid_detect_then_xrc_tidies_up_on_exception(
    mock_create_gridscan_callbacks: MagicMock,
    mock_get_ready_for_oav_and_close_shutter: MagicMock,
    run_engine: RunEngine,
    i04_grid_detect_then_xrc_default_params,
):
    mock_create_gridscan_callbacks.side_effect = CustomError
    with pytest.raises(CustomError):
        run_engine(
            i04_grid_detect_then_xrc_default_params(
                udc=False,
            )
        )

    assert mock_get_ready_for_oav_and_close_shutter.call_count == 1


@patch(
    "mx_bluesky.beamlines.i04.experiment_plans.i04_grid_detect_then_xray_centre_plan.get_ready_for_oav_and_close_shutter",
    autospec=True,
)
@patch(
    "mx_bluesky.beamlines.i04.experiment_plans.i04_grid_detect_then_xray_centre_plan.grid_detect_then_xray_centre",
    autospec=True,
)
@patch(
    "mx_bluesky.beamlines.i04.experiment_plans.i04_grid_detect_then_xray_centre_plan.create_gridscan_callbacks",
    autospec=True,
)
@patch(
    "mx_bluesky.beamlines.i04.experiment_plans.i04_grid_detect_then_xray_centre_plan.get_results_and_move_to_xtal",
    new=MagicMock(),
)
async def test_i04_grid_detect_then_xrc_sets_beamsize_before_grid_detect_then_reverts(
    mock_create_gridscan_callbacks: MagicMock,
    mock_grid_detect_then_xray_centre: MagicMock,
    mock_get_ready_for_oav_and_close_shutter: MagicMock,
    run_engine: RunEngine,
    i04_grid_detect_then_xrc_default_params: partial[MsgGenerator],
    transfocator: Transfocator,
    give_grid_common_specific_grid,
):
    initial_beamsize = 5.6
    set_mock_value(transfocator.current_vertical_size_rbv, initial_beamsize)
    parent_mock = MagicMock()
    assert isinstance(transfocator.set, MagicMock)
    parent_mock.attach_mock(transfocator.set, "transfocator_set")
    parent_mock.attach_mock(
        mock_create_gridscan_callbacks, "mock_create_gridscan_callbacks"
    )
    run_engine(i04_grid_detect_then_xrc_default_params())

    assert (
        mock_grid_detect_then_xray_centre.call_args.kwargs[
            "parameters"
        ].selected_aperture
        == ApertureValue.LARGE
    )
    assert parent_mock.method_calls == [
        call.transfocator_set(DEFAULT_XRC_BEAMSIZE_MICRONS),
        call.mock_create_gridscan_callbacks(),
        call.transfocator_set(initial_beamsize),
    ]


@patch(
    "mx_bluesky.beamlines.i04.experiment_plans.i04_grid_detect_then_xray_centre_plan.get_ready_for_oav_and_close_shutter",
    autospec=True,
)
@patch(
    "mx_bluesky.beamlines.i04.experiment_plans.i04_grid_detect_then_xray_centre_plan.grid_detect_then_xray_centre",
    autospec=True,
)
async def test_given_no_diffraction_found_i04_grid_detect_then_xrc_returns_sample_to_initial_position(
    mock_grid_detect_then_xray_centre: MagicMock,
    mock_get_ready_for_oav_and_close_shutter: MagicMock,
    run_engine: RunEngine,
    i04_grid_detect_then_xrc_default_params: partial[MsgGenerator],
    smargon: Smargon,
):
    initial_x, initial_y, initial_z = 1, 2, 3
    set_mock_value(smargon.x.user_readback, initial_x)
    set_mock_value(smargon.y.user_readback, initial_y)
    set_mock_value(smargon.z.user_readback, initial_z)

    mock_grid_detect_then_xray_centre.side_effect = CrystalNotFoundError

    with pytest.raises(CrystalNotFoundError):
        run_engine(i04_grid_detect_then_xrc_default_params())

    get_mock_put(smargon.x.user_setpoint).assert_has_calls([call(initial_x)])
    get_mock_put(smargon.y.user_setpoint).assert_has_calls([call(initial_y)])
    get_mock_put(smargon.z.user_setpoint).assert_has_calls([call(initial_z)])


@pytest.fixture()
def give_grid_common_specific_grid(test_fgs_params: SpecifiedThreeDGridScan):
    # Using this fixture means grid common always comes with its specific grid params
    class GridCommonWithSpecificGrid(GridCommon):
        @property
        def specified_grid_params(self):
            return test_fgs_params

    with patch(
        "mx_bluesky.beamlines.i04.experiment_plans.i04_grid_detect_then_xray_centre_plan.GridCommon",
        GridCommonWithSpecificGrid,
    ):
        yield


@patch(
    "mx_bluesky.beamlines.i04.experiment_plans.i04_grid_detect_then_xray_centre_plan.get_ready_for_oav_and_close_shutter",
    new=MagicMock(),
)
@patch(
    "mx_bluesky.beamlines.i04.experiment_plans.i04_grid_detect_then_xray_centre_plan.grid_detect_then_xray_centre",
    autospec=True,
)
@patch(
    "mx_bluesky.beamlines.i04.experiment_plans.i04_grid_detect_then_xray_centre_plan.setup_beamline_for_oav",
    autospec=True,
)
@patch(
    "mx_bluesky.beamlines.i04.experiment_plans.i04_grid_detect_then_xray_centre_plan.create_gridscan_callbacks",
    autospec=True,
)
@patch(
    "mx_bluesky.beamlines.i04.experiment_plans.i04_grid_detect_then_xray_centre_plan.fix_transmission_and_exposure_time_for_current_wavelength",
    return_value=(1, 0.004),
)
@patch(
    "mx_bluesky.beamlines.i04.experiment_plans.i04_grid_detect_then_xray_centre_plan.get_results_and_move_to_xtal",
    new=MagicMock(),
)
def test_i04_grid_detect_then_xrc_calculates_exposure_and_transmission_then_uses_grid_common(
    mock_fix_transmission: MagicMock,
    mock_create_gridscan_callbacks: MagicMock,
    mock_setup_beamline_for_oav: MagicMock,
    mock_grid_detect_then_xray_centre: MagicMock,
    i04_grid_detect_then_xrc_default_params: partial[MsgGenerator],
    run_engine: RunEngine,
    test_full_grid_scan_params: GridCommon,
    test_fgs_params: SpecifiedThreeDGridScan,
    give_grid_common_specific_grid,
):
    expected_trans_frac = 1
    expected_exposure_time = 0.004
    mock_fix_transmission.return_value = (expected_trans_frac, expected_exposure_time)

    run_engine(i04_grid_detect_then_xrc_default_params())
    mock_fix_transmission.assert_called_once()

    grid_common_params = mock_grid_detect_then_xray_centre.call_args.kwargs[
        "parameters"
    ]
    assert isinstance(grid_common_params, GridCommon)
    assert grid_common_params.exposure_time_s == expected_exposure_time
    assert grid_common_params.transmission_frac == expected_trans_frac


@patch(
    "mx_bluesky.beamlines.i04.experiment_plans.i04_grid_detect_then_xray_centre_plan.fix_transmission_and_exposure_time_for_current_wavelength",
)
def test_get_grid_common_params(
    mock_fix_trans_and_exposure: MagicMock,
    tmp_path,
):
    expected_trans_frac = 0.2
    expected_exposure_time = 0.007
    mock_fix_trans_and_exposure.return_value = (
        expected_trans_frac,
        expected_exposure_time,
    )
    entry_params = I04AutoXrcParams(
        sample_id=1,
        file_name="filename",
        visit="cm40607-5",
        detector_distance_mm=264.5,
        storage_directory=str(tmp_path),
    )
    grid_common_params = _get_grid_common_params(1, entry_params)
    assert grid_common_params.exposure_time_s == expected_exposure_time
    assert grid_common_params.transmission_frac == expected_trans_frac


@patch(
    "mx_bluesky.beamlines.i04.experiment_plans.i04_grid_detect_then_xray_centre_plan.grid_detect_then_xray_centre",
    new=MagicMock(),
)
@patch(
    "mx_bluesky.beamlines.i04.experiment_plans.i04_grid_detect_then_xray_centre_plan.fix_transmission_and_exposure_time_for_current_wavelength",
    new=MagicMock(return_value=(1, 0.004)),
)
@patch(
    "mx_bluesky.beamlines.i04.experiment_plans.i04_grid_detect_then_xray_centre_plan.get_results_and_move_to_xtal",
    new=MagicMock(),
)
def test_grid_detect_then_xrc_stages_and_unstages_zocalo_and_gets_results(
    sim_run_engine: RunEngineSimulator,
    i04_grid_detect_then_xrc_default_params: partial[MsgGenerator],
    give_grid_common_specific_grid,
):
    msgs = sim_run_engine.simulate_plan(i04_grid_detect_then_xrc_default_params())

    msgs = assert_message_and_return_remaining(
        msgs,
        predicate=lambda msg: msg.command == "stage"
        and msg.obj.name == "zocalo"
        and msg.kwargs["group"] == ZOCALO_STAGE_GROUP,
    )

    msgs = assert_message_and_return_remaining(
        msgs,
        predicate=lambda msg: msg.command == "unstage" and msg.obj.name == "zocalo",
    )


@patch(
    "mx_bluesky.common.experiment_plans.common_grid_detect_then_xray_centre_plan.ispyb_activation_wrapper",
    lambda plan, _: plan,
)
@patch(
    "mx_bluesky.common.experiment_plans.common_grid_detect_then_xray_centre_plan.grid_detection_plan",
    lambda *_: fake_generator(),
)
@patch(
    "mx_bluesky.common.experiment_plans.common_grid_detect_then_xray_centre_plan.common_flyscan_xray_centre",
    lambda *_: fake_generator(),
)
@patch(
    "mx_bluesky.common.experiment_plans.common_grid_detect_then_xray_centre_plan.GridDetectionCallback",
    new=MagicMock(),
)
@patch(
    "mx_bluesky.common.experiment_plans.common_grid_detect_then_xray_centre_plan.create_parameters_for_flyscan_xray_centre",
)
@patch(
    "mx_bluesky.beamlines.i04.experiment_plans.i04_grid_detect_then_xray_centre_plan.get_results_and_move_to_xtal"
)
def test_detect_grid_and_do_gridscan_gives_params_specified_grid(
    mock_change_aperture_then_move_to_xtal: MagicMock,
    mock_create_flyscan_params: MagicMock,
    test_fgs_params: SpecifiedThreeDGridScan,
    run_engine: RunEngine,
    i04_grid_detect_then_xrc_default_params: partial[MsgGenerator],
):
    mock_create_flyscan_params.return_value = test_fgs_params
    run_engine(
        i04_grid_detect_then_xrc_default_params(
            udc=False,
        )
    )
    mock_change_aperture_then_move_to_xtal.assert_called_once()
    assert mock_change_aperture_then_move_to_xtal.call_args[0][1] == test_fgs_params


@patch(
    "mx_bluesky.common.experiment_plans.common_grid_detect_then_xray_centre_plan.ispyb_activation_wrapper",
    lambda plan, _: plan,
)
@patch(
    "mx_bluesky.common.experiment_plans.common_grid_detect_then_xray_centre_plan.grid_detection_plan",
    lambda *_: fake_generator(),
)
@patch(
    "mx_bluesky.common.experiment_plans.common_grid_detect_then_xray_centre_plan.common_flyscan_xray_centre",
    lambda *_: fake_generator(),
)
@patch(
    "mx_bluesky.common.experiment_plans.common_grid_detect_then_xray_centre_plan.GridDetectionCallback",
    new=MagicMock(),
)
@patch(
    "mx_bluesky.common.experiment_plans.common_grid_detect_then_xray_centre_plan.create_parameters_for_flyscan_xray_centre",
)
@patch(
    "mx_bluesky.common.experiment_plans.change_aperture_then_move_plan._get_xrc_results",
)
def test_i04_grid_detect_then_xrc_only_sets_aperture_at_start_of_plan(
    mock_get_xrc_results: MagicMock,
    mock_create_flyscan_params: MagicMock,
    test_fgs_params: SpecifiedThreeDGridScan,
    run_engine: RunEngine,
    i04_grid_detect_then_xrc_default_params: partial[MsgGenerator],
):
    aperture_scatterguard: ApertureScatterguard = (
        i04_grid_detect_then_xrc_default_params.keywords["aperture_scatterguard"]
    )
    get_mock_put(aperture_scatterguard.selected_aperture).reset_mock()

    def fake_get_xrc_results_plan():
        results = [
            XRayCentreResult(
                centre_of_mass_mm=np.array([0.05, 0.15, 0.25]),
                bounding_box_mm=(
                    np.array([0.15, 0.15, 0.15]),
                    np.array([0.75, 0.75, 0.65]),
                ),
                max_count=105062,
                total_count=2387574,
                sample_id=12345,
            )
        ]
        yield from ()
        return results

    mock_get_xrc_results.return_value = fake_get_xrc_results_plan()

    mock_create_flyscan_params.return_value = test_fgs_params
    run_engine(
        i04_grid_detect_then_xrc_default_params(
            udc=False,
        )
    )

    # Out of beam for grid detect, large for gridscan, out again after tidy up
    assert get_mock_put(aperture_scatterguard.selected_aperture).call_args_list == [
        call(ApertureValue.OUT_OF_BEAM),
        call(ApertureValue.LARGE),
        call(ApertureValue.OUT_OF_BEAM),
    ]
