import pytest
from unittest.mock import MagicMock, patch
from aare.devices.aerotech import AerotechController, AEROTECH_HOME
from aare.common.beamline import MXBeamline
from aare.common.coordinate import AerotechCoordinate, Coordinate
from aare.common.exception_handler import AerotechCommunicationError


@pytest.fixture
def mock_aerotech_api():
    with patch('aarescan_client.ApiClient'), \
            patch('aarescan_client.DefaultApi') as mock_api_class, \
            patch('aarescan_client.Configuration'):
        mock_api = mock_api_class.return_value
        yield mock_api


@pytest.fixture
def aerotech_controller(mock_aerotech_api):
    with patch("aare.devices.aerotech.cfg_get", return_value="http://mx-x10sa-queue-01.psi.ch:5234"):
        controller = AerotechController(MXBeamline.X10SA)

    controller._AerotechController__api = mock_aerotech_api
    controller._AerotechController__simulated = False
    return controller


def test_init_x10sa(mock_aerotech_api):
    with patch("aare.devices.aerotech.cfg_get", return_value="http://mx-x10sa-queue-01.psi.ch:5234"):
        controller = AerotechController(MXBeamline.X10SA)
    assert controller._AerotechController__base == "http://mx-x10sa-queue-01.psi.ch:5234"
    assert controller._AerotechController__simulated is False


def test_init_simulated():
    controller = AerotechController(MXBeamline.SIMULATED)
    assert controller._AerotechController__simulated is True


def test_cancel(aerotech_controller, mock_aerotech_api):
    aerotech_controller.cancel()
    mock_aerotech_api.cancel_post.assert_called_once()


def test_is_idle(aerotech_controller, mock_aerotech_api):
    mock_aerotech_api.status_get.return_value.state = 'Idle'
    assert aerotech_controller.is_idle() is True

    mock_aerotech_api.status_get.return_value.state = 'Busy'
    assert aerotech_controller.is_idle() is False


def test_get_position(aerotech_controller, mock_aerotech_api):
    status = MagicMock()
    status.x.pos = 1.0
    status.y.pos = 2.0
    status.z.pos = 3.0
    status.u.pos = 90.0
    mock_aerotech_api.status_get.return_value = status

    pos = aerotech_controller.get_position()
    assert pos.at_mm.x == 1.0
    assert pos.at_mm.y == 2.0
    assert pos.at_mm.z == 3.0
    assert pos.omega_deg == 90.0


def test_move_home(aerotech_controller, mock_aerotech_api):
    aerotech_controller.move_home()
    mock_aerotech_api.position_post.assert_called_once()


def test_position(aerotech_controller, mock_aerotech_api):
    target = AerotechCoordinate(at_mm=Coordinate(x=1, y=2, z=3), omega_deg=45)
    aerotech_controller.position(target)
    mock_aerotech_api.position_post.assert_called_once()


def test_rotation_scan(aerotech_controller, mock_aerotech_api):
    aerotech_controller.rotation_scan(rotation_deg=360, time_sec=10, start_pos_deg=0)
    mock_aerotech_api.rotation_scan_post.assert_called_once()


def test_grid_scan(aerotech_controller, mock_aerotech_api):
    aerotech_controller.grid_scan(grid_elem_count_y=10, grid_elem_size_y_um=10, grid_elem_count_x=10,
                                  grid_elem_size_x_um=10, time_sec=5)
    mock_aerotech_api.grid_scan_post.assert_called_once()


def test_screening_scan(aerotech_controller, mock_aerotech_api):
    aerotech_controller.screening_scan(rotation_deg=10, wedge_deg=2, time_sec=1, steps=5)
    mock_aerotech_api.screening_post.assert_called_once()


def test_api_error(aerotech_controller, mock_aerotech_api):
    mock_aerotech_api.status_get.side_effect = Exception("API Error")
    with pytest.raises(AerotechCommunicationError):
        aerotech_controller.is_idle()