from __future__ import annotations

from functools import partial

import bluesky.plan_stubs as bps
import bluesky.preprocessors as bpp
from bluesky.utils import MsgGenerator
from dodal.common import inject
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,
    set_fast_grid_scan_params,
)
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.plans.preprocessors.verify_undulator_gap import (
    verify_undulator_gap_before_run_decorator,
)
from pydantic import BaseModel

from mx_bluesky.beamlines.i04.external_interaction.config_server import (
    get_i04_feature_settings,
)
from mx_bluesky.common.device_setup_plans.setup_zebra_and_shutter import (
    setup_zebra_for_gridscan,
    tidy_up_zebra_after_gridscan,
)
from mx_bluesky.common.experiment_plans.change_aperture_then_move_plan import (
    get_results_and_move_to_xtal,
)
from mx_bluesky.common.experiment_plans.common_flyscan_xray_centre_plan import (
    BeamlineSpecificFGSFeatures,
    construct_beamline_specific_fast_gridscan_features,
)
from mx_bluesky.common.experiment_plans.common_grid_detect_then_xray_centre_plan import (
    grid_detect_then_xray_centre,
)
from mx_bluesky.common.experiment_plans.inner_plans.xrc_results_utils import (
    zocalo_stage_decorator,
)
from mx_bluesky.common.experiment_plans.oav_snapshot_plan import (
    setup_beamline_for_oav,
)
from mx_bluesky.common.external_interaction.callbacks.common.zocalo_callback import (
    ZocaloCallback,
)
from mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback import (
    GridscanISPyBCallback,
    generate_start_info_from_omega_map,
)
from mx_bluesky.common.external_interaction.callbacks.xray_centre.nexus_callback import (
    GridscanNexusFileCallback,
)
from mx_bluesky.common.parameters.components import get_param_version
from mx_bluesky.common.parameters.constants import (
    EnvironmentConstants,
    OavConstants,
    PlanGroupCheckpointConstants,
    PlanNameConstants,
)
from mx_bluesky.common.parameters.device_composites import (
    GridDetectThenXRayCentreComposite,
)
from mx_bluesky.common.parameters.gridscan import (
    GridCommon,
    SpecifiedThreeDGridScan,
)
from mx_bluesky.common.preprocessors.preprocessors import (
    set_transmission_and_trigger_xbpm_feedback_before_collection_decorator,
)
from mx_bluesky.common.utils.exceptions import CrystalNotFoundError
from mx_bluesky.common.utils.log import LOGGER
from mx_bluesky.common.utils.utils import (
    fix_transmission_and_exposure_time_for_current_wavelength,
)
from mx_bluesky.common.utils.xrc_result import XRayCentreEventHandler

DEFAULT_XRC_BEAMSIZE_MICRONS = 20


class I04AutoXrcParams(BaseModel):
    sample_id: int
    file_name: str
    visit: str
    detector_distance_mm: float
    storage_directory: str


def _change_beamsize(
    transfocator: Transfocator, beamsize: float, parameters: GridCommon
):
    """i04 always uses the large aperture and changes beamsize with the transfocator.

    An aperture is needed to reduce scatter but the transfocator is best used for beamsize
    changes as it gives more flux compared to a bigger beam with a small aperture.
    """
    parameters.selected_aperture = ApertureValue.LARGE
    yield from bps.abs_set(
        transfocator, beamsize, group=PlanGroupCheckpointConstants.GRID_READY_FOR_DC
    )


# See https://github.com/DiamondLightSource/blueapi/issues/506 for using device composites
def i04_default_grid_detect_and_xray_centre(
    parameters: I04AutoXrcParams,
    aperture_scatterguard: ApertureScatterguard = inject("aperture_scatterguard"),
    attenuator: BinaryFilterAttenuator = inject("attenuator"),
    backlight: Backlight = inject("backlight"),
    beamstop: Beamstop = inject("beamstop"),
    beamsize: Beamsize = inject("beamsize"),
    dcm: DoubleCrystalMonochromator = inject("dcm"),
    zebra_fast_grid_scan: ZebraFastGridScanThreeD = inject("zebra_fast_grid_scan"),
    flux: Flux = inject("flux"),
    oav: OAV = inject("oav"),
    pin_tip_detection: PinTipDetection = inject("pin_tip_detection"),
    s4_slit_gaps: S4SlitGaps = inject("s4_slit_gaps"),
    undulator: UndulatorInKeV = inject("undulator"),
    xbpm_feedback: XBPMFeedback = inject("xbpm_feedback"),
    zebra: Zebra = inject("zebra"),
    robot: BartRobot = inject("robot"),
    sample_shutter: ZebraShutter = inject("sample_shutter"),
    eiger: EigerDetector = inject("eiger"),
    synchrotron: Synchrotron = inject("synchrotron"),
    zocalo: ZocaloResults = inject("zocalo"),
    smargon: Smargon = inject("gonio"),
    detector_motion: DetectorMotion = inject("detector_motion"),
    transfocator: Transfocator = inject("transfocator"),
    oav_config: str = OavConstants.OAV_CONFIG_JSON,
    udc: bool = False,
) -> MsgGenerator:
    """
    A composite plan which:
    - Uses the OAV to draw a virtual grid over the sample and to take snapshots of the sample
    - Scans through the grid to identify the crystal centre
    - Changes the aperture to match the beam size to the crystal size
    - Moves the sample to the crystal centre of mass

    i04's implementation of this plan is very similar to Hyperion. However, since i04
    isn't running in a continuous Bluesky UDC loop, we take additional steps in beamline
    tidy-up.
    """

    composite = GridDetectThenXRayCentreComposite(
        eiger,
        synchrotron,
        smargon,
        aperture_scatterguard,
        attenuator,
        zocalo,
        backlight,
        beamstop,
        beamsize,
        dcm,
        detector_motion,
        zebra_fast_grid_scan,
        flux,
        oav,
        pin_tip_detection,
        s4_slit_gaps,
        undulator,
        xbpm_feedback,
        zebra,
        robot,
        sample_shutter,
    )
    initial_beamsize = yield from bps.rd(transfocator.current_vertical_size_rbv)

    initial_x = yield from bps.rd(smargon.x.user_readback)
    initial_y = yield from bps.rd(smargon.y.user_readback)
    initial_z = yield from bps.rd(smargon.z.user_readback)

    _current_wavelength_a = yield from bps.rd(composite.dcm.wavelength_in_a)
    grid_common_params = _get_grid_common_params(_current_wavelength_a, parameters)

    def tidy_beamline():
        yield from bps.mv(transfocator, initial_beamsize)

        if not udc:
            yield from get_ready_for_oav_and_close_shutter(
                composite.gonio,
                composite.backlight,
                composite.aperture_scatterguard,
                composite.detector_motion,
            )

    @zocalo_stage_decorator(composite.zocalo)
    @bpp.finalize_decorator(tidy_beamline)
    def _inner_grid_detect_then_xrc():
        # These callbacks let us talk to ISPyB and Nexgen. They aren't included in the common plan because
        # Hyperion handles its callbacks differently to BlueAPI-managed plans, see
        # https://github.com/DiamondLightSource/mx-bluesky/issues/1117
        flyscan_event_handler = XRayCentreEventHandler()
        callbacks = *create_gridscan_callbacks(), flyscan_event_handler

        @bpp.subs_decorator(callbacks)
        @verify_undulator_gap_before_run_decorator(composite)
        @set_transmission_and_trigger_xbpm_feedback_before_collection_decorator(
            composite,
            grid_common_params.transmission_frac,
            PlanNameConstants.GRIDSCAN_OUTER,
        )
        def grid_detect_then_xray_centre_with_callbacks():
            yield from grid_detect_then_xray_centre(
                composite=composite,
                parameters=grid_common_params,
                xrc_params_type=SpecifiedThreeDGridScan,
                construct_beamline_specific=construct_i04_specific_features,
                oav_config=oav_config,
            )

        try:
            yield from grid_detect_then_xray_centre_with_callbacks()
            assert isinstance(
                grid_common_params.specified_grid_params, SpecifiedThreeDGridScan
            ), "Specified grid params couldn't be found after grid detection"
            yield from get_results_and_move_to_xtal(
                composite,
                grid_common_params.specified_grid_params,
                flyscan_event_handler,
            )
        except CrystalNotFoundError:
            yield from bps.mv(
                smargon.x, initial_x, smargon.y, initial_y, smargon.z, initial_z
            )
            raise

    yield from _change_beamsize(
        transfocator, DEFAULT_XRC_BEAMSIZE_MICRONS, grid_common_params
    )
    yield from _inner_grid_detect_then_xrc()


def get_ready_for_oav_and_close_shutter(
    smargon: Smargon,
    backlight: Backlight,
    aperture_scatterguard: ApertureScatterguard,
    detector_motion: DetectorMotion,
):
    yield from bps.wait(PlanGroupCheckpointConstants.GRID_READY_FOR_DC)
    group = "get_ready_for_oav_and_close_shutter"
    LOGGER.info("Non-udc tidy: Setting up beamline for OAV")
    yield from setup_beamline_for_oav(
        smargon, backlight, aperture_scatterguard, group=group
    )
    LOGGER.info("Non-udc tidy: Closing detector shutter")
    yield from bps.abs_set(
        detector_motion.shutter,
        0,
        group=group,
    )
    yield from bps.wait(group)


def create_gridscan_callbacks() -> tuple[
    GridscanNexusFileCallback, GridscanISPyBCallback
]:
    return (
        GridscanNexusFileCallback(param_type=SpecifiedThreeDGridScan),
        GridscanISPyBCallback(
            param_type=GridCommon,
            emit=ZocaloCallback(
                PlanNameConstants.DO_FGS,
                EnvironmentConstants.ZOCALO_ENV,
                generate_start_info_from_omega_map,
            ),
        ),
    )


def construct_i04_specific_features(
    xrc_composite: GridDetectThenXRayCentreComposite,
    xrc_parameters: SpecifiedThreeDGridScan,
) -> BeamlineSpecificFGSFeatures:
    """
    Get all the information needed to do the i04 XRC flyscan.
    """
    signals_to_read_pre_flyscan = [
        xrc_composite.undulator.current_gap,
        xrc_composite.synchrotron.synchrotron_mode,
        xrc_composite.s4_slit_gaps,
        xrc_composite.gonio,
        xrc_composite.dcm.energy_in_keV,
    ]

    signals_to_read_during_collection = [
        xrc_composite.aperture_scatterguard,
        xrc_composite.attenuator.actual_transmission,
        xrc_composite.flux.flux_reading,
        xrc_composite.dcm.energy_in_keV,
        xrc_composite.eiger.bit_depth,
        xrc_composite.beamsize,
        xrc_composite.eiger.cam.roi_mode,
        xrc_composite.eiger.ispyb_detector_id,
    ]

    tidy_plan = partial(
        tidy_up_zebra_after_gridscan,
        xrc_composite.zebra,
        xrc_composite.sample_shutter,
        group="flyscan_zebra_tidy",
        wait=True,
    )
    set_flyscan_params_plan = partial(
        set_fast_grid_scan_params,
        xrc_composite.zebra_fast_grid_scan,
        xrc_parameters.fast_gridscan_params,
    )
    fgs_motors = xrc_composite.zebra_fast_grid_scan
    return construct_beamline_specific_fast_gridscan_features(
        partial(
            setup_zebra_for_gridscan,
        ),
        tidy_plan,
        set_flyscan_params_plan,
        fgs_motors,
        signals_to_read_pre_flyscan,
        signals_to_read_during_collection,  # type: ignore # until https://github.com/DiamondLightSource/mx-bluesky/issues/1076
    )


def _get_grid_common_params(
    _current_wavelength_a: float, parameters: I04AutoXrcParams
) -> GridCommon:
    """Calculate scaled transmission and exposure by comparing current beamline energy to default energy"""
    feature_settings = get_i04_feature_settings()
    _assumed_wavelength_a = feature_settings.ASSUMED_WAVELENGTH_IN_A
    _unscaled_transmission = feature_settings.XRC_UNSCALED_TRANSMISSION_FRAC
    _unscaled_exposure_time_s = feature_settings.XRC_UNSCALED_EXPOSURE_TIME_S
    transmission_frac, exposure_time_s = (
        fix_transmission_and_exposure_time_for_current_wavelength(
            _current_wavelength_a,
            _assumed_wavelength_a,
            _unscaled_transmission,
            _unscaled_exposure_time_s,
        )
    )

    return GridCommon(
        sample_id=parameters.sample_id,
        file_name=parameters.file_name,
        visit=parameters.visit,
        detector_distance_mm=parameters.detector_distance_mm,
        storage_directory=parameters.storage_directory,
        transmission_frac=transmission_frac,
        exposure_time_s=exposure_time_s,
        parameter_model_version=get_param_version(),
    )
