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.logger_config import setup_logger
from aaredaqlib.models import StagePositionEnum, SampleCameraSettings, ZoomModeEnum
from aaredevices.mx_lib import pv_wait

logger = setup_logger("aareDAQ")

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:
            logger.info("DTZ still moving.")
        raise RuntimeError("Timeout moving devs. Dropping to maintenance.")


def wait_for_se_devices(devs: BeamlineDevices, cfg: BeamlineConfig):
    hub = cfg.settings
    hub_cryo = cfg.cryojet_settings

    _cryo = hub_cryo.cryojet_in_use
    _cryopark = hub_cryo.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):
    logger.info(" moving COLLIMATOR to PARK position")
    devs.collimator = StagePositionEnum.PARK
    logger.info(" moving SMARGON to HOME position")
    devs.smargon.move_home(wait=True)
    logger.info(" moving AEROTECH to MOUNT position")
    devs.aerotech.move(ABR_OMEGA_MOUNT, wait=True, direct=True)
    devs.abr_pos = ABR_POS_MOUNT
    logger.info(" moving REFLECTOR to DOWN position")
    devs.reflector_up = False
    logger.info(" LOCKING AEROTERCH")
    devs.aerotech.lock()
    logger.info(f" Smargon position in RSE: {devs.smargon.readback} ABS: {devs.abr_pos.x} {devs.abr_pos.y} {devs.abr_pos.z}")


def m2se(devs: BeamlineDevices, cfg: BeamlineConfig):
    print("executing MAINTENANCE -> sample exchange")
    devs.detector_cover.put(1) # Close detector cover

    hub = cfg.settings
    hub_cryo = cfg.cryojet_settings

    cfg.zoom_mode = ZoomModeEnum.User
    devs.samcam_settings = cfg.zoom_settings.get_camera_settings(devs.zoom)

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

    _dtzpark = hub.dtz_park
    _cryopark = hub_cryo.cryojet_park_position
    _cryo = hub_cryo.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
    hub_cryo = cfg.cryojet_settings
    logger.info(" moving cryojet to park position")
    _cryopark = hub_cryo.cryojet_park_position
    _dtzpark = hub.dtz_park

    devs.aerotech.move(ABR_OMEGA_MOUNT)

    if devs.cryo.position < _cryopark:
        devs.cryo.move(_cryopark)
    logger.info(" moving COLLIMATOR to DOWN position")
    devs.collimator = StagePositionEnum.DOWN
    logger.info(" moving REFLECTOR to DOWN position")
    devs.reflector_up = False
    devs.beamstop_stage_up = False
    logger.info(" moving SMARGON to HOME position")
    devs.smargon.move_home(wait=True)
    logger.info(" moving DETECTOR COVER to DOWN position")
    devs.detector_cover.put(1) # Close detector cover
    if devs.dtz.position < _dtzpark:
        logger.info(" moving DETECTOR to PARK position")
        devs.dtz.move(_dtzpark)


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


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


def se2sa(devs: BeamlineDevices, cfg: BeamlineConfig):
    hub = cfg.settings
    cfg.zoom_mode = ZoomModeEnum.User
    devs.detector_cover.put(1)
    devs.samcam_settings = cfg.zoom_settings.get_camera_settings(devs.zoom)
    hub_cryo = cfg.cryojet_settings
    _cryo = hub_cryo.cryojet_in_use
    cryo_meas = hub_cryo.cryojet_measurement_position
    _dtz_safety = hub.dtz_bsz_safety_margin
    devs.aerotech.unlock()
    logger.info("aerotech unlocked")
    devs.aerotech.set_direct_mode()
    logger.info("aerotech in direct mode")
    logger.info(" moving reflector to up position")
    devs.reflector_up = True
    logger.info("moving beamstop to up position")
    devs.beamstop_stage_up = True
    logger.info(" setting lamp to 2.5")
    devs.lamp_light = 2.5
    logger.info(" moving aerotech to measure position")
    devs.abr_pos = cfg.abr_meas_pos

    if _cryo:
        logger.info(" moving cryojet to measure position")
        devs.cryo.move(cryo_meas)
    logger.info("moving detector to measure position")
    devs.dtz.move(cfg.dtz, wait=False)

def rse2sa(devs: BeamlineDevices, cfg: BeamlineConfig):
    se2sa(devs, cfg)

def sa2dc(devs: BeamlineDevices, cfg: BeamlineConfig):
    devs.detector_cover.put(2)  # Open detector cover

    devs.reflector_up = False

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

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


def dc2sa(devs: BeamlineDevices, cfg: BeamlineConfig):
    hub = cfg.settings
    devs.detector_cover.put(1) # Close detector cover
    devs.collimator = StagePositionEnum.PARK
    devs.reflector_up = True

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


def sa2xrf(devs: BeamlineDevices, cfg: BeamlineConfig):
    """sample alignment to XrfCollection"""
    pass


def xrf2sa(devs: BeamlineDevices, cfg: BeamlineConfig):
    """XrfCollection to sample alignment"""
    pass


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
    hub_cryo = cfg.cryojet_settings

    cryo_meas = hub_cryo.cryojet_measurement_position
    _cryo = hub_cryo.cryojet_in_use

    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
    hub_cryo = cfg.cryojet_settings
    _cryopark = hub_cryo.cryojet_park_position
    _cryo = hub_cryo.cryojet_in_use

    if _cryo:
        devs.cryo.move(_cryopark)

    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):
    val=1000
    devs.zoom = val
    cfg.zoom_mode = ZoomModeEnum.BeamLocation
    devs.detector_cover.put(1) # Close detector cover
    devs.samcam_settings = cfg.zoom_settings.get_camera_settings(val)
    devs.lamp_light = 2.5
    devs.abr_pos = Coordinate(x=-26, y=0, z=0)
    devs.cryo.move(80)
    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
    hub_cryo = cfg.cryojet_settings
    cryo_meas = hub_cryo.cryojet_measurement_position
    _cryo = hub_cryo.cryojet_in_use

    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

    val=1
    devs.zoom = val
    cfg.zoom_mode = ZoomModeEnum.User
    devs.samcam_settings = cfg.zoom_settings.get_camera_settings(val)

    if _cryo:
        devs.cryo.move(cryo_meas)

    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):
    logger.info("moved to dewer transfer")
    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)
    logger.info("return to sample alignment")