import pytest
from aarecommon.models.models import DewarAddress, SampleShortInfo
from PySide6.QtCore import Qt

from aare.gui.models.sample_queue_model import SampleQueueSpreadsheet
from aare.gui.models.user_sample_model import UserSampleSpreadsheet


@pytest.fixture
def sample_list():
    return [
        SampleShortInfo(
            db_id=1,
            puck_name="P1",
            dewar_name="D1",
            sample_name="S1",
            run_number=1,
            user="U1",
            pin=1,
            location=DewarAddress(segment="A", pos=1),
        ),
        SampleShortInfo(
            db_id=2,
            puck_name="P2",
            dewar_name="D2",
            sample_name="S2",
            run_number=2,
            user="U2",
            pin=2,
            location=DewarAddress(segment="A", pos=2),
        ),
        SampleShortInfo(
            db_id=3,
            puck_name="P1",
            dewar_name="D1",
            sample_name="S3",
            run_number=3,
            user="U1",
            pin=3,
            location=DewarAddress(segment="A", pos=1),
        ),
    ]


def test_user_sample_model_init(sample_list):
    model = UserSampleSpreadsheet(samples=sample_list)
    model.set_show_all_pgroups(True)  # Ensure all pgroups are shown for testing
    assert model.rowCount() == 3
    # Check if filtering works
    # "User" is column 5
    model.set_filter("User", "U2")
    assert model.rowCount() == 1
    model.clear_filter()
    assert model.rowCount() == 3


def test_user_sample_model_column_filter(sample_list):
    model = UserSampleSpreadsheet(samples=sample_list)
    model.set_show_all_pgroups(True)
    # Column 5 is user
    model.set_column_filter(5, "U1")
    assert model.rowCount() == 2
    model.clear_all_column_filters()
    assert model.rowCount() == 3


def test_user_sample_model_unique_values(sample_list):
    model = UserSampleSpreadsheet(samples=sample_list)
    model.set_show_all_pgroups(True)
    # Column 5 is User
    users = model.unique_values_for_column(5)
    assert "U1" in users
    assert "U2" in users
    assert len(users) == 2


def test_sample_queue_model_init(sample_list):
    model = SampleQueueSpreadsheet(samples=sample_list[:2])
    assert model.rowCount() == 2
    assert model.columnCount() == 3
    assert model.data(model.index(0, 0), Qt.ItemDataRole.DisplayRole) == "D1"
    assert model.data(model.index(1, 2), Qt.ItemDataRole.DisplayRole) == "S2"


def test_sample_queue_model_update(sample_list):
    model = SampleQueueSpreadsheet()
    assert model.rowCount() == 0
    model.updateData(sample_list[:2])
    assert model.rowCount() == 2


def test_sample_queue_model_remove(sample_list):
    model = SampleQueueSpreadsheet(samples=sample_list[:2])
    model.remove_sample(1)
    assert model.rowCount() == 1
    assert model.samples[0].db_id == 2


def test_sample_queue_model_clear(sample_list):
    model = SampleQueueSpreadsheet(samples=sample_list[:2])
    model.clearSamples()
    assert model.rowCount() == 0


def test_sample_queue_model_set_running(sample_list):
    model = SampleQueueSpreadsheet(samples=sample_list[:2])
    # Background color role for first row
    color_not_running = model.data(model.index(0, 0), Qt.ItemDataRole.BackgroundRole)
    model.set_running(True)
    color_running = model.data(model.index(0, 0), Qt.ItemDataRole.BackgroundRole)
    assert color_not_running != color_running


def test_sample_queue_model_header(sample_list):
    model = SampleQueueSpreadsheet(samples=sample_list[:2])
    assert model.headerData(0, Qt.Orientation.Horizontal, Qt.ItemDataRole.DisplayRole) == "Dewar"
    assert model.headerData(1, Qt.Orientation.Vertical, Qt.ItemDataRole.DisplayRole) == "2"


def test_sample_queue_model_flags(sample_list):
    model = SampleQueueSpreadsheet(samples=sample_list[:2])
    flags = model.flags(model.index(0, 0))
    assert flags & Qt.ItemFlag.ItemIsDropEnabled