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

import numpy as np
import pytest
from bluesky.simulators import RunEngineSimulator
from bluesky.utils import Msg
from dodal.devices.beamsize.beamsize import BeamsizeBase
from dodal.devices.synchrotron import SynchrotronMode
from dodal.devices.zocalo import ZocaloResults
from ophyd_async.core import AsyncStatus, completed_status, set_mock_value

from mx_bluesky.common.experiment_plans.common_flyscan_xray_centre_plan import (
    BeamlineSpecificFGSFeatures,
)
from mx_bluesky.common.external_interaction.ispyb.ispyb_store import (
    IspybIds,
    StoreInIspyb,
)
from mx_bluesky.common.utils.xrc_result import XRayCentreResult
from mx_bluesky.hyperion.experiment_plans.hyperion_flyscan_xray_centre_plan import (
    construct_hyperion_specific_features,
)
from mx_bluesky.hyperion.experiment_plans.robot_load_and_change_energy import (
    RobotLoadAndEnergyChangeComposite,
)
from mx_bluesky.hyperion.experiment_plans.robot_load_then_centre_plan import (
    RobotLoadThenCentreComposite,
)
from mx_bluesky.hyperion.external_interaction.callbacks.__main__ import (
    create_gridscan_callbacks,
)
from mx_bluesky.hyperion.parameters.device_composites import (
    HyperionFlyScanXRayCentreComposite,
)
from mx_bluesky.hyperion.parameters.gridscan import HyperionSpecifiedThreeDGridScan

FLYSCAN_RESULT_HIGH = XRayCentreResult(
    centre_of_mass_mm=np.array([0.1, 0.2, 0.3]),
    bounding_box_mm=(np.array([0.09, 0.19, 0.29]), np.array([0.11, 0.21, 0.31])),
    max_count=30,
    total_count=100,
    sample_id=2,
)
FLYSCAN_RESULT_MED = XRayCentreResult(
    centre_of_mass_mm=np.array([0.4, 0.5, 0.6]),
    bounding_box_mm=(np.array([0.09, 0.19, 0.29]), np.array([0.11, 0.21, 0.31])),
    max_count=20,
    total_count=120,
    sample_id=1,
)
FLYSCAN_RESULT_LOW = XRayCentreResult(
    centre_of_mass_mm=np.array([0.7, 0.8, 0.9]),
    bounding_box_mm=(np.array([0.09, 0.19, 0.29]), np.array([0.11, 0.21, 0.31])),
    max_count=10,
    total_count=140,
    sample_id=1,
)
FLYSCAN_RESULT_HIGH_NO_SAMPLE_ID = XRayCentreResult(
    centre_of_mass_mm=np.array([0.1, 0.2, 0.3]),
    bounding_box_mm=(np.array([0.09, 0.19, 0.29]), np.array([0.11, 0.21, 0.31])),
    max_count=30,
    total_count=100,
    sample_id=None,
)


@pytest.fixture(autouse=True)
def always_use_i03_beamline(use_beamline_i03): ...


@pytest.fixture
def sim_run_engine_for_rotation(sim_run_engine):
    sim_run_engine.add_handler(
        "read",
        lambda msg: {"values": {"value": SynchrotronMode.USER}},
        "synchrotron-synchrotron_mode",
    )
    sim_run_engine.add_handler(
        "read",
        lambda msg: {"values": {"value": -1}},
        "synchrotron-top_up_start_countdown",
    )
    sim_run_engine.add_handler(
        "read", lambda msg: {"values": {"value": -1}}, "gonio_omega"
    )
    return sim_run_engine


def mock_zocalo_trigger(zocalo: ZocaloResults, result):
    @AsyncStatus.wrap
    async def mock_complete(results):
        await zocalo._put_results(results, {"dcid": 0, "dcgid": 0})

    zocalo.trigger = MagicMock(side_effect=partial(mock_complete, result))


def modified_store_grid_scan_mock(*args, dcids=(0, 0), dcgid=0, **kwargs):
    mock = MagicMock(spec=StoreInIspyb)
    mock.begin_deposition.return_value = IspybIds(
        data_collection_ids=dcids, data_collection_group_id=dcgid
    )
    mock.update_deposition.return_value = IspybIds(
        data_collection_ids=dcids, data_collection_group_id=dcgid, grid_ids=(0, 0)
    )
    return mock


@pytest.fixture
def mock_subscriptions():
    with (
        patch(
            "mx_bluesky.common.external_interaction.callbacks.common.zocalo_callback.ZocaloTrigger",
            autospec=True,
        ),
        patch(
            "mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback.StoreInIspyb.append_to_comment"
        ),
        patch(
            "mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback.StoreInIspyb.begin_deposition",
            new=MagicMock(
                return_value=IspybIds(
                    data_collection_ids=(0, 0), data_collection_group_id=0
                )
            ),
        ),
        patch(
            "mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback.StoreInIspyb.update_deposition",
            new=MagicMock(
                return_value=IspybIds(
                    data_collection_ids=(0, 0),
                    data_collection_group_id=0,
                    grid_ids=(0, 0),
                )
            ),
        ),
    ):
        nexus_callback, ispyb_callback = create_gridscan_callbacks()
        ispyb_callback.ispyb = MagicMock(spec=StoreInIspyb)

    return (nexus_callback, ispyb_callback)


def fake_read(obj, initial_positions, _):
    initial_positions[obj] = 0
    yield Msg("null", obj)


@pytest.fixture
def robot_load_composite(
    smargon,
    dcm,
    robot,
    aperture_scatterguard,
    oav,
    webcam,
    thawer,
    lower_gonio,
    eiger,
    xbpm_feedback,
    attenuator,
    beamstop_phase1,
    fast_grid_scan,
    undulator,
    undulator_dcm,
    s4_slit_gaps,
    vfm,
    mirror_voltages,
    backlight,
    detector_motion,
    flux,
    pin_tip_detection_with_found_pin,
    zocalo,
    synchrotron,
    sample_shutter,
    zebra,
    panda,
    panda_fast_grid_scan,
    beamsize: BeamsizeBase,
) -> RobotLoadThenCentreComposite:
    set_mock_value(dcm.energy_in_keV.user_readback, 11.105)
    smargon.stub_offsets.set = MagicMock(side_effect=lambda _: completed_status())
    aperture_scatterguard.set = MagicMock(side_effect=lambda _: completed_status())
    set_mock_value(smargon.omega.max_velocity, 131)
    return RobotLoadThenCentreComposite(
        xbpm_feedback=xbpm_feedback,
        attenuator=attenuator,
        aperture_scatterguard=aperture_scatterguard,
        backlight=backlight,
        beamsize=beamsize,
        beamstop=beamstop_phase1,
        detector_motion=detector_motion,
        eiger=eiger,
        zebra_fast_grid_scan=fast_grid_scan,
        flux=flux,
        oav=oav,
        pin_tip_detection=pin_tip_detection_with_found_pin,
        gonio=smargon,
        synchrotron=synchrotron,
        s4_slit_gaps=s4_slit_gaps,
        undulator=undulator,
        zebra=zebra,
        zocalo=zocalo,
        panda=panda,
        panda_fast_grid_scan=panda_fast_grid_scan,
        thawer=thawer,
        sample_shutter=sample_shutter,
        vfm=vfm,
        mirror_voltages=mirror_voltages,
        dcm=dcm,
        undulator_dcm=undulator_dcm,
        robot=robot,
        webcam=webcam,
        lower_gonio=lower_gonio,
    )


@pytest.fixture
def robot_load_and_energy_change_composite(
    smargon,
    dcm,
    robot,
    aperture_scatterguard,
    oav,
    webcam,
    thawer,
    lower_gonio,
    vfm,
    mirror_voltages,
    undulator_dcm,
    xbpm_feedback,
    attenuator,
    backlight,
) -> RobotLoadAndEnergyChangeComposite:
    composite = RobotLoadAndEnergyChangeComposite(
        vfm,
        mirror_voltages,
        dcm,
        undulator_dcm,
        xbpm_feedback,
        attenuator,
        robot,
        webcam,
        lower_gonio,
        thawer,
        oav,
        smargon,
        aperture_scatterguard,
        backlight,
    )
    composite.gonio.stub_offsets.set = MagicMock(
        side_effect=lambda _: completed_status()
    )
    composite.aperture_scatterguard.set = MagicMock(
        side_effect=lambda _: completed_status()
    )
    set_mock_value(composite.dcm.energy_in_keV.user_readback, 11.105)

    return composite


def sim_fire_event_on_open_run(sim_run_engine: RunEngineSimulator, run_name: str):
    def fire_event(msg: Msg):
        try:
            sim_run_engine.fire_callback("start", msg.kwargs)
        except Exception as e:
            print(f"Exception is {e}")

    def msg_maches_run(msg: Msg):
        return msg.run == run_name

    sim_run_engine.add_handler("open_run", fire_event, msg_maches_run)


@pytest.fixture
def grid_detection_callback_with_detected_grid():
    with patch(
        "mx_bluesky.common.experiment_plans.common_grid_detect_then_xray_centre_plan.GridDetectionCallback",
        autospec=True,
    ) as callback:
        callback.return_value.get_grid_parameters.return_value = {
            "transmission_frac": 1.0,
            "exposure_time_s": 0,
            "x_start_um": 0,
            "y_start_um": 0,
            "y2_start_um": 0,
            "z_start_um": 0,
            "z2_start_um": 0,
            "x_steps": 10,
            "y_steps": 10,
            "z_steps": 10,
            "x_step_size_um": 0.1,
            "y_step_size_um": 0.1,
            "z_step_size_um": 0.1,
        }
        yield callback


@pytest.fixture
def beamline_specific(
    hyperion_flyscan_xrc_composite: HyperionFlyScanXRayCentreComposite,
    hyperion_fgs_params: HyperionSpecifiedThreeDGridScan,
) -> BeamlineSpecificFGSFeatures:
    return construct_hyperion_specific_features(
        hyperion_flyscan_xrc_composite, hyperion_fgs_params
    )
