from types import SimpleNamespace
from unittest.mock import MagicMock

from aare.daq.operations.rotation.service import RotationService


def _img(number, angle=0.0, spots=0, res=0.0):
    return SimpleNamespace(number=number, angle=angle, spots=spots, res=res)


def _make_service(jfjoch, aare):
    ctx = SimpleNamespace(deps=SimpleNamespace(jfjoch=jfjoch, aare=aare))
    return RotationService(context=ctx, logger=MagicMock())


def test_screening_diffraction_uploads_one_image_per_wedge():
    jfjoch = MagicMock()
    jfjoch.get_diffraction_image.side_effect = lambda image_id, **kw: f"jpeg-{image_id}".encode()
    aare = MagicMock()
    svc = _make_service(jfjoch, aare)

    scan_result = SimpleNamespace(images=[_img(0, 0.0, 245, 2.1), _img(1, 90.0, 12, 8.5)])
    svc._ingest_screening_diffraction(SimpleNamespace(db_id=123), scan_result)

    # one fetch + one upload per wedge, with spots + resolution-estimate ring
    # + beam center overlays on
    assert jfjoch.get_diffraction_image.call_count == 2
    for call in jfjoch.get_diffraction_image.call_args_list:
        assert call.kwargs["show_spots"] is True
        assert call.kwargs["show_res_est"] is True
        assert call.kwargs["show_beam_center"] is True

    assert aare.upload_jpg.call_count == 2
    first = aare.upload_jpg.call_args_list[0]
    assert first.args[:3] == (123, "123_screening_diffraction_0", b"jpeg-0")
    assert "Screening diffraction (wedge 1" in first.kwargs["message"]
    assert "245 spots" in first.kwargs["message"]


def test_screening_diffraction_skips_images_without_id():
    jfjoch = MagicMock()
    jfjoch.get_diffraction_image.return_value = b"jpeg"
    aare = MagicMock()
    svc = _make_service(jfjoch, aare)

    svc._ingest_screening_diffraction(
        SimpleNamespace(db_id=7), SimpleNamespace(images=[SimpleNamespace(number=None), _img(3)])
    )

    jfjoch.get_diffraction_image.assert_called_once_with(
        3, show_spots=True, show_res_est=True, show_beam_center=True
    )
    assert aare.upload_jpg.call_count == 1


def test_screening_diffraction_tolerates_fetch_failure():
    jfjoch = MagicMock()
    jfjoch.get_diffraction_image.side_effect = [Exception("buffer miss"), b"jpeg-6"]
    aare = MagicMock()
    svc = _make_service(jfjoch, aare)

    svc._ingest_screening_diffraction(
        SimpleNamespace(db_id=9), SimpleNamespace(images=[_img(5), _img(6)])
    )

    # first wedge failed and was skipped; second still uploaded
    assert aare.upload_jpg.call_count == 1
    assert aare.upload_jpg.call_args.args[1] == "9_screening_diffraction_6"