"""
This module contains the bluesky plan entry points for use with hyperion.
hyperion-blueapi loads this module via blueapi.py which contains module-load initialisation
hooks that are run when blueapi loads the plan module, which are not needed by the
in-process runner and so not defined here.

The json schema and documentation therein generated by the blueapi /plans endpoint
from this file constitutes the hyperion-blueapi interface to the hyperion supervisor
process.
"""

from bluesky import plan_stubs as bps
from bluesky.utils import MsgGenerator
from dodal.common import inject
from dodal.devices.aperturescatterguard import ApertureScatterguard
from dodal.devices.detector.detector_motion import DetectorMotion, ShutterState
from dodal.devices.motors import XYZStage
from dodal.devices.robot import BartRobot
from dodal.devices.smargon import Smargon

from mx_bluesky.common.device_setup_plans.robot_load_unload import (
    robot_unload as _robot_unload,
)
from mx_bluesky.hyperion.blueapi.parameters import (
    LoadCentreCollectParams,
    load_centre_collect_to_internal,
)
from mx_bluesky.hyperion.experiment_plans.load_centre_collect_full_plan import (
    LoadCentreCollectComposite,
)
from mx_bluesky.hyperion.experiment_plans.load_centre_collect_full_plan import (
    load_centre_collect_full as _load_centre_collect_full,
)
from mx_bluesky.hyperion.experiment_plans.udc_default_state import (
    UDCDefaultDevices,
)
from mx_bluesky.hyperion.experiment_plans.udc_default_state import (
    move_to_udc_default_state as _move_to_udc_default_state,
)

__all__ = [
    "LoadCentreCollectComposite",
    "LoadCentreCollectParams",
    "UDCDefaultDevices",
    "clean_up_udc",
    "load_centre_collect",
    "move_to_udc_default_state",
    "robot_unload",
]


def load_centre_collect(
    parameters: LoadCentreCollectParams,
    composite: LoadCentreCollectComposite = inject(),
) -> MsgGenerator:
    """
    Attempt a complete data collection experiment, consisting of the following:
        * Load the sample if necessary
        * Move to the specified goniometer start angles
        * Perform optical centring, then X-ray centring
        * If X-ray centring finds one or more diffracting centres then for each centre
    that satisfies the chosen selection function,
    move to that centre and do a collection with the specified parameters.
    """
    yield from _load_centre_collect_full(
        composite, load_centre_collect_to_internal(parameters)
    )


def robot_unload(
    visit: str,
    robot: BartRobot = inject("robot"),
    smargon: Smargon = inject("gonio"),
    aperture_scatterguard: ApertureScatterguard = inject("aperture_scatterguard"),
    lower_gonio: XYZStage = inject("lower_gonio"),
) -> MsgGenerator:
    """
    Unload the currently mounted pin into the location that it was loaded from.
    This is to be invoked as the final step upon successful completion of the UDC queue.
    """
    yield from _robot_unload(robot, smargon, aperture_scatterguard, lower_gonio, visit)


def clean_up_udc(
    cleanup_group: str = "cleanup",
    detector_motion: DetectorMotion = inject("detector_motion"),
) -> MsgGenerator:
    """
    Perform actions that are always performed on ending UDC, regardless of whether
    it terminated abnormally or not.
    Currently the only thing this does is close the detector shutter.
    Args:
        detector_motion:

    Returns:
    """
    yield from bps.abs_set(
        detector_motion.shutter, ShutterState.CLOSED, group=cleanup_group
    )
    yield from bps.wait(cleanup_group)


def move_to_udc_default_state(
    composite: UDCDefaultDevices = inject(),
) -> MsgGenerator:
    """
    Move beamline hardware to known positions prior to UDC start.
    """
    yield from _move_to_udc_default_state(composite)
