import pytest
from aarecommon.models.models import (
    BeamlineStateEnum,
    BeamMarkCoeffModel,
    DataCollectionParameters,
    DewarAddress,
    MLBoxType,
    MLOutputModel,
    SampleShortInfo,
)


def test_sample_short_info_methods():
    info = SampleShortInfo(
        db_id=1,
        puck_name="puck1",
        dewar_name="dew1",
        sample_name="sample1",
        run_number=1,
        aaredb_params=DataCollectionParameters(
            totalangle=180, processingresolution=1.5, cloud=True
        ),
        user="group1",
        pin=3,
        location=DewarAddress(segment="A", pos=2),
    )

    addr = info.tell_address()
    assert addr.puck.segment == "A"
    assert addr.puck.pos == 2
    assert addr.pin == 3

    assert info.loc_str() == "A2-3"
    assert info.loc_str_sort() == "A2-03"
    assert info.aaredb_params is not None
    assert info.aaredb_params.totalangle == 180
    assert info.aaredb_params.processingresolution == 1.5

    info_no_loc = info.model_copy(update={"location": None})
    assert info_no_loc.loc_str() == "-"
    assert info_no_loc.loc_str_sort() == ""

    data = {
        "db_id": 2,
        "puck_name": "puck2",
        "dewar_name": "dew2",
        "sample_name": "sample2",
        "run_number": 2,
        "aaredb_params": {"totalrange": 120, "userresolution": 1.8, "cloud": ""},
        "user": "group2",
        "pin": 4,
        "location": {"segment": "B", "pos": 5},
    }
    info2 = SampleShortInfo.from_dict(data)
    assert info2.db_id == 2
    assert info2.location.segment == "B"
    assert info2.aaredb_params is not None
    assert info2.aaredb_params.totalangle == 120
    assert info2.aaredb_params.processingresolution == 1.8
    assert info2.aaredb_params.cloud is True


def test_beam_mark_coeff_model_apply():
    model = BeamMarkCoeffModel(coeff_x=(1.0, 2.0, 5.0), coeff_y=(3.0, 4.0, 6.0))
    res = model.apply(10.0)
    assert res.x == 125.0
    assert res.y == 346.0


def test_ml_output_model_extra_methods():
    model = MLOutputModel()
    key = model.add_box(MLBoxType.CRYSTAL, (1, 2, 3, 4), 0.8)

    box_model = model.get_box_model(key)
    assert box_model.conf == 0.8

    assert model.get_box_tuple(key) == (1, 2, 3, 4)
    assert model.get_box_tuple("NonExistent") is None

    assert model.get_box_tuple_with_conf(key) == (1, 2, 3, 4, 0.8)
    assert model.get_box_tuple_with_conf("NonExistent") is None

    tuples = model.get_tuples_for_class(MLBoxType.CRYSTAL)
    assert len(tuples) == 1
    assert tuples[0] == (1, 2, 3, 4)

    tuples_conf = model.get_tuples_with_conf_for_class(MLBoxType.CRYSTAL)
    assert len(tuples_conf) == 1
    assert tuples_conf[0] == (1, 2, 3, 4, 0.8)

    assert MLOutputModel.get_class_str(MLBoxType.LOOP_ALL) == "Loop_all"
    assert MLOutputModel.get_class_str(MLBoxType.PIN) == "Pin"
    assert MLOutputModel.get_class_str(MLBoxType.CRYSTAL) == "Crystal"
    assert MLOutputModel.get_class_str(MLBoxType.LOOP_FACE) == "Loop_face"
    assert MLOutputModel.get_class_str(MLBoxType.ICE) == "Ice"
    assert MLOutputModel.get_class_str(MLBoxType.NEEDLE) == "Needle"
    assert MLOutputModel.get_class_str(100) == "Unknown"


def test_beamline_state_enum_display_name():
    assert BeamlineStateEnum.SampleExchange.display_name() == "Sample exchange"
    assert BeamlineStateEnum.Moving.display_name() == "Moving"
    assert BeamlineStateEnum.display_name(None) == "-"