import pytest
import json
from unittest.mock import MagicMock, patch
from aare.daq.spreadsheetupdater import on_message, get_ws_headers, set_spreadsheet_in_redis
from aare.common.models import SampleShortInfoList


@pytest.fixture
def mock_config():
    with patch('aare.daq.spreadsheetupdater.config') as mock:
        mock._BeamlineConfig__bl = "X10SA"
        mock._BeamlineConfig__client = MagicMock()
        # Mocking private attributes access which the code uses
        mock._BeamlineConfig__client = mock._BeamlineConfig__client
        mock._BeamlineConfig__bl = mock._BeamlineConfig__bl
        yield mock


def test_get_ws_headers_success():
    with patch('os.getenv', return_value="secret"):
        headers = get_ws_headers()
        assert headers == ["X-Shared-Password: secret"]


def test_get_ws_headers_fail():
    with patch('os.getenv', return_value=None):
        with pytest.raises(ValueError):
            get_ws_headers()


def test_set_spreadsheet_in_redis(mock_config):
    data = {"test": "data"}
    with patch('aare.daq.spreadsheetupdater.config') as mock_cfg_internal:
        mock_client = MagicMock()
        mock_cfg_internal._BeamlineConfig__client = mock_client
        mock_cfg_internal.client = mock_client

        set_spreadsheet_in_redis(data)

        found = False
        for attr in dir(mock_cfg_internal):
            val = getattr(mock_cfg_internal, attr)
            if isinstance(val, MagicMock) and val.set.called:
                found = True
                break
        assert found or mock_client.set.called


def test_on_message_success(mock_config):
    message = json.dumps({
        "samples": [
            {
                "id": 1,
                "puck_name": "P1",
                "puck_type": "UniPuck",
                "puck_location_in_dewar": 1,
                "dewar_id": 1,
                "pgroup": "p12345",
                "dewar_name": "D1",
                "tell_position": "A1",
                "samples": [
                    {
                        "id": 1,
                        "sample_name": "S1",
                        "run_number": 100,
                        "pgroup": "p12345",
                        "position": 1,
                        "priority": 1,
                        "mount_count": 0,
                        "rotation_count": 0,
                        "raster_count": 0,
                        "screening_count": 0,
                        "data_collection_parameters": {}
                    }
                ]
            },
            {
                "id": 2,
                "puck_name": "Ref",
                "puck_type": "UniPuck",
                "puck_location_in_dewar": 2,
                "dewar_id": 1,
                "pgroup": "p12345",
                "dewar_name": "D1",
                "tell_position": "X1",
                "samples": [
                    {
                        "id": 2,
                        "sample_name": "R1",
                        "run_number": 1,
                        "pgroup": "p12345",
                        "position": 1,
                        "priority": 1,
                        "mount_count": 0,
                        "rotation_count": 0,
                        "raster_count": 0,
                        "screening_count": 0,
                        "data_collection_parameters": {}
                    }
                ]
            }
        ]
    })

    on_message(None, message)

    normal_key = "X10SA:sample_spreadsheet"
    calls = mock_config._BeamlineConfig__client.set.call_args_list
    assert any(call.args[0] == normal_key for call in calls)

    ref_key = "X10SA:reference-tools"
    assert any(call.args[0] == ref_key for call in calls)


def test_on_message_empty_ref(mock_config):
    message = json.dumps({
        "samples": [
            {
                "id": 1,
                "puck_name": "P1",
                "puck_type": "UniPuck",
                "puck_location_in_dewar": 1,
                "dewar_id": 1,
                "pgroup": "p12345",
                "dewar_name": "D1",
                "tell_position": "A1",
                "samples": []
            }
        ]
    })

    on_message(None, message)

    ref_key = "X10SA:reference-tools"
    mock_config._BeamlineConfig__client.delete.assert_called_with(ref_key)


def test_on_message_invalid_json(mock_config):
    on_message(None, "invalid json")
    mock_config._BeamlineConfig__client.set.assert_not_called()