import time

from epics import poll

from aaredaq.config import ABR_POS_MOUNT, ABR_OMEGA_MOUNT
from aaredaq.devices import BeamlineDevices
from aaredaq.config import BeamlineConfig
from aaredaqlib.coordinate import Coordinate
from aaredaqlib.models import StagePositionEnum, SampleCameraSettings
from mxlibs3.mx_lib import pv_wait


def move_bsz(devs: BeamlineDevices, target: float):
    if abs(target - devs.bsz.position) > 0.1:
        beamstop_stage_measure = devs.beamstop_stage_up
        reflector_measure = devs.reflector_up
        devs.reflector_up = False
        devs.beamstop_stage_up = True

        devs.bsz.move(target, wait=True)

        if reflector_measure:
            devs.reflector_up = True

        if not beamstop_stage_measure:
            devs.beamstop_stage_up = False


def wait_for_dc_devices(devs: BeamlineDevices):
    timeout = time.time() + 60
    timeisup = False
    # devs.detector_cover.wait()
    while not timeisup and (
            not devs.dtz.done_moving
    ):
        poll(0.1)
        timeisup = timeout < time.time()
    if timeisup:
        if not devs.dtz.done_moving:
            print("DTZ still moving.")
        raise RuntimeError("Timeout moving devs. Dropping to maintenance.")


def wait_for_se_devices(devs: BeamlineDevices, cfg: BeamlineConfig):
    hub = cfg.settings
    _cryo = hub.cryojet_in_use
    _cryopark = hub.cryojet_park_position

    if _cryo and not devs.cryo.position_is(_cryopark):
        raise RuntimeError("Cryojet did not reach far position")


def common_2rse(devs: BeamlineDevices, cfg: BeamlineConfig):
    devs.collimator_up = False
    devs.smargon.move_home(wait=True)
    devs.aerotech.move(ABR_OMEGA_MOUNT, wait=True, direct=True)
    devs.abr_pos = ABR_POS_MOUNT
    devs.aerotech.lock()


def m2se(devs: BeamlineDevices, cfg: BeamlineConfig):
    print("executing MAINTENANCE -> sample exchange")
    hub = cfg.settings

    devs.collimator = StagePositionEnum.PARK
    devs.scintillator = StagePositionEnum.DOWN

    _dtzpark = hub.dtz_park
    _cryopark = hub.cryojet_park_position
    _cryo = hub.cryojet_in_use

    # devs.detector_cover.move("closed")
    if not devs.aerotech.is_ready_and_willing():
        devs.aerotech.stop()
        time.sleep(5.0)
        raise RuntimeError("Aerotech device is not ready")

    if devs.dtz.position < _dtzpark:
        devs.dtz.move(_dtzpark)

    devs.gmx.set_speed(100.0)
    devs.abr_pos = ABR_POS_MOUNT
    devs.aerotech.move(ABR_OMEGA_MOUNT)

    if _cryo:
        devs.cryo.move(_cryopark)

    devs.reflector_up = False

    devs.bsx.move(0.0)
    devs.bsy.move(0.0)
    move_bsz(devs, hub.bsz)

    devs.beamstop_stage_up = False


def sa2se(devs: BeamlineDevices, cfg: BeamlineConfig):
    hub = cfg.settings
    _cryopark = hub.cryojet_park_position
    _dtzpark = hub.dtz_park

    if devs.cryo.position < _cryopark:
        devs.cryo.move(_cryopark)

    devs.collimator_up = False
    devs.reflector_up = False
    devs.beamstop_stage_up = False
    devs.smargon.move_home(wait=True)
    if devs.dtz.position < _dtzpark:
        devs.dtz.move(_dtzpark)


def sa2rse(devs: BeamlineDevices, cfg: BeamlineConfig):
    common_2rse(devs, cfg)


def dc2rse(devs: BeamlineDevices, cfg: BeamlineConfig):
    if not devs.aerotech.is_ready_and_willing():
        devs.aerotech.stop()
        time.sleep(5.0)
        raise RuntimeError("Aerotech device is not ready")
    common_2rse(devs, cfg)


def se2sa(devs: BeamlineDevices, cfg: BeamlineConfig):
    hub = cfg.settings
    _cryo = hub.cryojet_in_use
    cryo_meas = hub.cryojet_measurement_position
    _dtz_safety = hub.dtz_bsz_safety_margin

    devs.aerotech.unlock()
    devs.aerotech.set_direct_mode()
    devs.reflector_up = True
    devs.beamstop_stage_up = True
    devs.lamp_light = 2.5
    devs.abr_pos = cfg.abr_meas_pos

    if _cryo:
        devs.cryo.move(cryo_meas)

    devs.dtz.move(cfg.dtz)


def sa2dc(devs: BeamlineDevices, cfg: BeamlineConfig):
    hub = cfg.settings

    _cryo = hub.cryojet_in_use
    cryo_meas = hub.cryojet_measurement_position

    devs.reflector_up = False

    if _cryo:
        devs.cryo.move(cryo_meas)

    #devs.reflector.wait("park")
    devs.beamstop_stage_up = True
    devs.dtz.move(cfg.dtz, wait=True)
    devs.collimator_up = True

    pv_wait(devs.dtz, None, tolerance=1.0)
    wait_for_dc_devices(devs)


def dc2se(devs: BeamlineDevices, cfg: BeamlineConfig):
    hub = cfg.settings
    _cryo = hub.cryojet_in_use
    _dtzpark = hub.dtz_park

    devs.cryo.move(hub.cryojet_park_position)
    devs.aerotech.move(ABR_OMEGA_MOUNT, speed=180, wait=True, direct=True)

    # devs.detector_cover.move("closed")

    if not devs.aerotech.is_ready_and_willing():
        devs.aerotech.stop()
        time.sleep(5.0)
        raise RuntimeError("Aerotech device is not ready")

    if devs.dtz.position < _dtzpark:
        devs.dtz.move(_dtzpark)
    devs.collimator_up = False
    devs.reflector_up = True

def dc2sa(devs: BeamlineDevices, cfg: BeamlineConfig):
    hub = cfg.settings
    devs.collimator_up = False
    devs.reflector_up = True

    # devs.gmz.move(cfg.abr_meas_pos.z, wait=True)


def sa2xrf(devs: BeamlineDevices, cfg: BeamlineConfig):
    """sample alignment to XrfCollection"""
    if devs.dtz.position < 150.0:
        devs.dtz.move(150.0)
    devs.reflector_up = False
    # devs.xrf_stage.move("measure")
    devs.beamstop_stage_up = True
    move_bsz(devs, 70.0)


def xrf2sa(devs: BeamlineDevices, cfg: BeamlineConfig):
    """XrfCollection to sample alignment"""
    if devs.dtz.position < 150.0:
        devs.dtz.move(150.0)
    devs.reflector_up = False
    # devs.xrf_stage.move("measure")
    devs.beamstop_stage_up = True
    move_bsz(devs, 70.0)


def sa2ws(devs: BeamlineDevices, cfg: BeamlineConfig):
    hub = cfg.settings
    _dtzwash = hub.dtz_wash_sample_distance
    devs.cryo.move(hub.cryojet_park_position)
    if devs.dtz.position < _dtzwash:
        devs.dtz.move(_dtzwash)
    devs.reflector_up = False
    devs.beamstop_stage_up = False


def ws2sa(devs: BeamlineDevices, cfg: BeamlineConfig):
    hub = cfg.settings
    cryo_meas = hub.cryojet_measurement_position
    cryo_user = hub.cryojet_measurement_user
    _cryo = hub.cryojet_in_use

    if cryo_user >= cryo_meas:
        cryo_meas = cryo_user
    devs.dtz.move(cfg.dtz)
    if _cryo:
        devs.cryo.move(cryo_meas)
    devs.reflector_up = True
    devs.beamstop_stage_up = True


def sa2ba(devs: BeamlineDevices, cfg: BeamlineConfig):
    hub = cfg.settings

    if hub.cryojet_in_use:
        devs.cryo.move(hub.cryojet_park_position)
    devs.gmx.move(-38.0)
    devs.lamp_light = 2.5
    devs.beamstop_stage_up = True
    move_bsz(devs, 0.0)


def ba2sa(devs: BeamlineDevices, cfg: BeamlineConfig):
    hub = cfg.settings
    devs.lamp_light = 2.5
    devs.beamstop_stage_up = True
    move_bsz(devs, hub.bsz)
    devs.reflector_up = True
    devs.abr_pos = cfg.abr_meas_pos


def sa2bl(devs: BeamlineDevices, cfg: BeamlineConfig):
    devs.zoom = 1000
    devs.samcam_settings = SampleCameraSettings(exposure=0.02, gain=100)
    devs.lamp_light = 2.5

    devs.abr_pos = cfg.abr_meas_pos + Coordinate(x=-10)
    devs.cryo.move(20)
    devs.set_scintillator(StagePositionEnum.MEASURE, wait=False)
    devs.set_collimator(StagePositionEnum.MEASURE, wait=False)
    devs.beamstop_stage_up = True
    devs.reflector_up = False
    devs.scintillator = StagePositionEnum.MEASURE
    devs.collimator = StagePositionEnum.MEASURE
    devs.shutter = True

def bl2sa(devs: BeamlineDevices, cfg: BeamlineConfig):
    hub = cfg.settings

    devs.shutter = False

    devs.set_scintillator(StagePositionEnum.DOWN, wait=False)
    devs.set_collimator(StagePositionEnum.PARK, wait=False)

    devs.reflector_up = True
    devs.scintillator = StagePositionEnum.DOWN
    devs.collimator = StagePositionEnum.PARK

    if hub.cryojet_in_use:
        devs.cryo.move(hub.cryojet_measurement_position)

    devs.abr_pos = cfg.abr_meas_pos


def bl2ba(devs: BeamlineDevices, cfg: BeamlineConfig):
    """beam location to beamstop alignment"""
    devs.gmx.move(-38)
    devs.reflector_up = True
    # devs.shutter.move("closed")
    time.sleep(0.5)
    move_bsz(devs, 0)


def ba2bl(devs: BeamlineDevices, cfg: BeamlineConfig):
    hub = cfg.settings
    move_bsz(devs, hub.bsz)
    devs.abr_pos = cfg.abr_meas_pos + Coordinate(x=1)
    devs.reflector_up = True


def sa2dh(devs: BeamlineDevices, cfg: BeamlineConfig):
    if devs.tell.get_mounted_sample() is not None:
        try:
            # Best effort try to unmount
            devs.tell.unmount(wait=True)
        except Exception as e:
            print(f"Error for unmounting: {e}")
    devs.tell.dry(wait_cold=-1, wait=False)


def dh2sa(devs: BeamlineDevices, cfg: BeamlineConfig):
    devs.tell.move_cold(wait=True)