from __future__ import annotations

import dataclasses
from collections.abc import Generator
from datetime import datetime
from pathlib import Path
from typing import cast

import bluesky.plan_stubs as bps
import bluesky.preprocessors as bpp
from blueapi.core import BlueskyContext
from bluesky.utils import Msg
from dodal.devices.aperturescatterguard import ApertureScatterguard, ApertureValue
from dodal.devices.attenuator import Attenuator
from dodal.devices.dcm import DCM
from dodal.devices.focusing_mirror import FocusingMirrorWithStripes, VFMMirrorVoltages
from dodal.devices.motors import XYZPositioner
from dodal.devices.oav.oav_detector import OAV
from dodal.devices.robot import BartRobot, SampleLocation
from dodal.devices.smargon import Smargon, StubPosition
from dodal.devices.thawer import Thawer
from dodal.devices.undulator_dcm import UndulatorDCM
from dodal.devices.webcam import Webcam
from dodal.devices.xbpm_feedback import XBPMFeedback
from dodal.plans.motor_util_plans import MoveTooLarge, home_and_reset_wrapper

from mx_bluesky.hyperion.experiment_plans.set_energy_plan import (
    SetEnergyComposite,
    set_energy_plan,
)
from mx_bluesky.hyperion.log import LOGGER
from mx_bluesky.hyperion.parameters.constants import CONST
from mx_bluesky.hyperion.parameters.robot_load import RobotLoadAndEnergyChange


@dataclasses.dataclass
class RobotLoadAndEnergyChangeComposite:
    # SetEnergyComposite fields
    vfm: FocusingMirrorWithStripes
    vfm_mirror_voltages: VFMMirrorVoltages
    dcm: DCM
    undulator_dcm: UndulatorDCM
    xbpm_feedback: XBPMFeedback
    attenuator: Attenuator

    # RobotLoad fields
    robot: BartRobot
    webcam: Webcam
    lower_gonio: XYZPositioner
    thawer: Thawer
    oav: OAV
    smargon: Smargon
    aperture_scatterguard: ApertureScatterguard


def create_devices(context: BlueskyContext) -> RobotLoadAndEnergyChangeComposite:
    from mx_bluesky.hyperion.utils.context import device_composite_from_context

    return device_composite_from_context(context, RobotLoadAndEnergyChangeComposite)


def wait_for_smargon_not_disabled(smargon: Smargon, timeout=60):
    """Waits for the smargon disabled flag to go low. The robot hardware is responsible
    for setting this to low when it is safe to move. It does this through a physical
    connection between the robot and the smargon.
    """
    LOGGER.info("Waiting for smargon enabled")
    SLEEP_PER_CHECK = 0.1
    times_to_check = int(timeout / SLEEP_PER_CHECK)
    for _ in range(times_to_check):
        smargon_disabled = yield from bps.rd(smargon.disabled)
        if not smargon_disabled:
            LOGGER.info("Smargon now enabled")
            return
        yield from bps.sleep(SLEEP_PER_CHECK)
    raise TimeoutError(
        "Timed out waiting for smargon to become enabled after robot load"
    )


def take_robot_snapshots(oav: OAV, webcam: Webcam, directory: Path):
    time_now = datetime.now()
    snapshot_format = f"{time_now.strftime('%H%M%S')}_{{device}}_after_load"
    for device in [oav.snapshot, webcam]:
        yield from bps.abs_set(
            device.filename, snapshot_format.format(device=device.name)
        )
        yield from bps.abs_set(device.directory, str(directory))
        # Note: should be able to use `wait=True` after https://github.com/bluesky/bluesky/issues/1795
        yield from bps.trigger(device, group="snapshots")
        yield from bps.wait("snapshots")


def prepare_for_robot_load(
    aperture_scatterguard: ApertureScatterguard, smargon: Smargon
):
    yield from bps.abs_set(
        aperture_scatterguard,
        ApertureValue.ROBOT_LOAD,
        group="prepare_robot_load",
    )

    yield from bps.mv(smargon.stub_offsets, StubPosition.RESET_TO_ROBOT_LOAD)

    # fmt: off
    yield from bps.mv(smargon.x, 0,
                      smargon.y, 0,
                      smargon.z, 0,
                      smargon.omega, 0,
                      smargon.chi, 0,
                      smargon.phi, 0)
    # fmt: on

    yield from bps.wait("prepare_robot_load")


def do_robot_load(
    composite: RobotLoadAndEnergyChangeComposite,
    sample_location: SampleLocation,
    demand_energy_ev: float | None,
    thawing_time: float,
):
    yield from bps.abs_set(
        composite.robot,
        sample_location,
        group="robot_load",
    )

    if demand_energy_ev:
        yield from set_energy_plan(
            demand_energy_ev / 1000,
            cast(SetEnergyComposite, composite),
        )

    yield from bps.wait("robot_load")

    yield from bps.abs_set(
        composite.thawer.thaw_for_time_s, thawing_time, group="thawing_finished"
    )
    yield from wait_for_smargon_not_disabled(composite.smargon)


def raise_exception_if_moved_out_of_cryojet(exception):
    yield from bps.null()
    if isinstance(exception, MoveTooLarge):
        raise Exception(
            f"Moving {exception.axis} back to {exception.position} after \
                        robot load would move it out of the cryojet. The max safe \
                        distance is {exception.maximum_move}"
        )


def pin_already_loaded(
    robot: BartRobot, sample_location: SampleLocation
) -> Generator[Msg, None, bool]:
    current_puck = yield from bps.rd(robot.current_puck)
    current_pin = yield from bps.rd(robot.current_pin)
    return (
        int(current_puck) == sample_location.puck
        and int(current_pin) == sample_location.pin
    )


def robot_load_and_snapshots(
    composite: RobotLoadAndEnergyChangeComposite,
    location: SampleLocation,
    snapshot_directory: Path,
    thawing_time: float,
    demand_energy_ev: float | None,
):
    robot_load_plan = do_robot_load(
        composite,
        location,
        demand_energy_ev,
        thawing_time,
    )

    # The lower gonio must be in the correct position for the robot load and we
    # want to put it back afterwards. Note we don't wait the robot is interlocked
    # to the lower gonio and the  move is quicker than the robot takes to get to the
    # load position.
    yield from bpp.contingency_wrapper(
        home_and_reset_wrapper(
            robot_load_plan,
            composite.lower_gonio,
            BartRobot.LOAD_TOLERANCE_MM,
            CONST.HARDWARE.CRYOJET_MARGIN_MM,
            "lower_gonio",
            wait_for_all=False,
        ),
        except_plan=raise_exception_if_moved_out_of_cryojet,
    )

    yield from take_robot_snapshots(composite.oav, composite.webcam, snapshot_directory)

    yield from bps.create(name=CONST.DESCRIPTORS.ROBOT_LOAD)
    yield from bps.read(composite.robot.barcode)
    yield from bps.read(composite.oav.snapshot)
    yield from bps.read(composite.webcam)
    yield from bps.save()

    yield from bps.wait("reset-lower_gonio")


def robot_load_and_change_energy_plan(
    composite: RobotLoadAndEnergyChangeComposite,
    params: RobotLoadAndEnergyChange,
):
    assert params.sample_puck is not None
    assert params.sample_pin is not None

    sample_location = SampleLocation(params.sample_puck, params.sample_pin)

    yield from prepare_for_robot_load(
        composite.aperture_scatterguard, composite.smargon
    )
    yield from bpp.run_wrapper(
        robot_load_and_snapshots(
            composite,
            sample_location,
            params.snapshot_directory,
            params.thawing_time,
            params.demand_energy_ev,
        ),
        md={
            "subplan_name": CONST.PLAN.ROBOT_LOAD,
            "metadata": {
                "visit": params.visit,
                "sample_id": params.sample_id,
                "sample_puck": sample_location.puck,
                "sample_pin": sample_location.pin,
            },
            "activate_callbacks": [
                "RobotLoadISPyBCallback",
            ],
        },
    )
