import pytest
from unittest.mock import MagicMock, patch
from aare.daq.workflows import move_bsz, common_2rse, sa2se, sa2rse, sa2xtal_snapshot, dc2xtal_snapshot, xtal_snapshot2dc, xtal_snapshot2sa, dc2rse, se2sa, sa2dc, dc2sa, sa2xrf, sa2dh, dh2sa
from aare.daq.config import ABR_POS_MOUNT
from aare.common.models import StagePositionEnum
from aare.devices.area_detector import AutoEnum
from aare.devices.bec_worker import BeamlineState

@pytest.fixture
def mock_devs():
    devs = MagicMock()
    devs.bsz.position = 0.0
    devs.beamstop_stage_up = False
    devs.reflector_up = False
    devs.bec_worker = None
    return devs

@pytest.fixture
def mock_cfg():
    cfg = MagicMock()
    cfg.abr_meas_pos = "MEAS_POS"
    return cfg

def test_move_bsz_no_move(mock_devs):
    mock_devs.bsz.position = 1.0
    move_bsz(mock_devs, 1.05)
    # No changes should happen as diff < 0.1
    assert mock_devs.beamstop_z.call_count == 0

def test_move_bsz_with_move(mock_devs):
    mock_devs.bsz.position = 0.0
    mock_devs.beamstop_stage_up = False
    mock_devs.reflector_up = True
    
    move_bsz(mock_devs, 1.0)
    
    # Check if reflector was toggled
    assert mock_devs.reflector_up == True
    # Check if beamstop stage was toggled
    assert mock_devs.beamstop_stage_up == False
    # Check if beamstop_z was set
    assert mock_devs.beamstop_z == 1.0

def test_common_2rse_no_bec(mock_devs, mock_cfg):
    mock_devs.bec_worker = None
    common_2rse(mock_devs, mock_cfg)
    
    assert mock_devs.collimator == 20.0
    assert mock_devs.scintillator == 20.0
    mock_devs.smargon_move_home.assert_called_once()
    assert mock_devs.beamstop_stage_up == StagePositionEnum.PARK

def test_common_2rse_with_bec(mock_devs, mock_cfg):
    mock_devs.bec_worker = MagicMock()
    common_2rse(mock_devs, mock_cfg)
    
    mock_devs.bec_worker.planner.move_to.assert_called_with(BeamlineState.ROBOT_SAMPLE_EXCHANGE)
    mock_devs.smargon_move_home.assert_called_once()
    assert mock_devs.aerotech_pos == ABR_POS_MOUNT

def test_sa2se(mock_devs, mock_cfg):
    mock_devs.bec_worker = MagicMock()
    sa2se(mock_devs, mock_cfg)
    
    mock_devs.samcam_auto.assert_any_call(AutoEnum.AUTO)
    mock_devs.bec_worker.planner.move_to.assert_called_with(BeamlineState.MANUAL_SAMPLE_EXCHANGE)
    assert mock_devs.aerotech_pos == mock_cfg.abr_meas_pos
    mock_devs.smargon_move_home.assert_called_once()
    mock_devs.samcam_auto.assert_any_call(AutoEnum.ONCE)

def test_sa2rse(mock_devs, mock_cfg):
    sa2rse(mock_devs, mock_cfg)
    mock_devs.samcam_auto.assert_any_call(AutoEnum.AUTO)
    mock_devs.samcam_auto.assert_any_call(AutoEnum.ONCE)

def test_sa2xtal_snapshot(mock_devs, mock_cfg):
    mock_devs.bec_worker = MagicMock()
    sa2xtal_snapshot(mock_devs, mock_cfg)
    assert mock_devs.samcam_settings.exposure == 0.001
    mock_devs.bec_worker.planner.move_to.assert_called_with(BeamlineState.XTAL_SNAPSHOT)

def test_dc2xtal_snapshot(mock_devs, mock_cfg):
    mock_devs.bec_worker = MagicMock()
    dc2xtal_snapshot(mock_devs, mock_cfg)
    assert mock_devs.samcam_settings.exposure == 0.001
    mock_devs.bec_worker.planner.move_to.assert_called_with(BeamlineState.XTAL_SNAPSHOT)

def test_xtal_snapshot2dc(mock_devs, mock_cfg):
    mock_devs.bec_worker = MagicMock()
    xtal_snapshot2dc(mock_devs, mock_cfg)
    mock_devs.bec_worker.planner.move_to.assert_called_with(BeamlineState.DATA_COLLECTION)

def test_xtal_snapshot2sa(mock_devs, mock_cfg):
    mock_devs.bec_worker = MagicMock()
    xtal_snapshot2sa(mock_devs, mock_cfg)
    mock_devs.bec_worker.planner.move_to.assert_called_with(BeamlineState.SAMPLE_ALIGNMENT)

def test_dc2rse(mock_devs, mock_cfg):
    dc2rse(mock_devs, mock_cfg)
    mock_devs.samcam_auto.assert_any_call(AutoEnum.AUTO)

def test_se2sa(mock_devs, mock_cfg):
    mock_devs.bec_worker = MagicMock()
    se2sa(mock_devs, mock_cfg)
    mock_devs.bec_worker.planner.move_to.assert_called_with(BeamlineState.SAMPLE_ALIGNMENT)
    assert mock_devs.aerotech_pos == mock_cfg.abr_meas_pos

def test_sa2dc(mock_devs, mock_cfg):
    mock_devs.bec_worker = MagicMock()
    sa2dc(mock_devs, mock_cfg)
    mock_devs.bec_worker.move_to.assert_called_with(BeamlineState.DATA_COLLECTION)

def test_dc2sa(mock_devs, mock_cfg):
    mock_devs.bec_worker = MagicMock()
    dc2sa(mock_devs, mock_cfg)
    mock_devs.bec_worker.planner.move_to.assert_called_with(BeamlineState.SAMPLE_ALIGNMENT)
    assert mock_devs.aerotech_pos == mock_cfg.abr_meas_pos

def test_sa2xrf(mock_devs, mock_cfg):
    mock_devs.bec_worker = MagicMock()
    sa2xrf(mock_devs, mock_cfg)
    mock_devs.bec_worker.planner.move_to.assert_called_with(BeamlineState.FLUX_MEASUREMENT)

def test_sa2dh(mock_devs, mock_cfg):
    mock_devs.tell.get_mounted_sample.return_value = "Sample"
    sa2dh(mock_devs, mock_cfg)
    mock_devs.tell.unmount.assert_called_with(wait=True)
    mock_devs.tell.dry.assert_called_with(wait_cold=-1, wait=False)

def test_dh2sa(mock_devs, mock_cfg):
    mock_devs.bec_worker = MagicMock()
    dh2sa(mock_devs, mock_cfg)
    mock_devs.bec_worker.planner.move_to.assert_called_with(BeamlineState.SAMPLE_ALIGNMENT)