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 6 is User (column 0 is the frozen #+status cell)
    model.set_column_filter(6, "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 6 is User (column 0 is the frozen #+status cell)
    users = model.unique_values_for_column(6)
    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


# --- Status logic of the combined dewar/queue view ---------------------------
# The dewar table doubles as the queue view: the frozen "#" column carries a
# status fill (mounted > queued > flagged > measured) and the chip row filters
# by status. This is the logic a local contact trusts at a glance, so it gets
# its own tests.


def _status(model, row):
    brush = model.data(model.index(row, 0), Qt.ItemDataRole.BackgroundRole)
    return None if brush is None else brush.color().name().lower()


def _row_of(model, db_id):
    return next(r for r in range(model.rowCount()) if model.get_id(r).db_id == db_id)


@pytest.fixture
def status_model(sample_list):
    from aarecommon.models.models import DewarAddress, SampleShortInfo

    # A measured sample: any rotation data counts (exactly 1 MUST count;
    # unmeasured is rotation_count 0, like the fixture's samples 1-3).
    sample_list.append(
        SampleShortInfo(
            db_id=4,
            puck_name="P3",
            dewar_name="D3",
            sample_name="S4",
            run_number=4,
            user="U1",
            pin=4,
            rotation_count=1,
            location=DewarAddress(segment="B", pos=1),
        )
    )
    model = UserSampleSpreadsheet(samples=sample_list)
    model.set_show_all_pgroups(True)
    return model


def test_status_color_priority(status_model):
    from aare.gui.styles import (
        SAMPLE_ROW_QUEUED_BG,
        SAMPLE_STATUS_FLAGGED_BG,
        SAMPLE_STATUS_MEASURED_BG,
        SAMPLE_STATUS_QUEUED_BG,
    )

    model = status_model
    assert _status(model, _row_of(model, 1)) is None

    model.set_queued_ids({1})
    model.set_flagged(1, True)
    # Queued beats flagged in the All view.
    assert _status(model, _row_of(model, 1)) == SAMPLE_STATUS_QUEUED_BG.lower()
    model.set_queued_ids(set())
    assert _status(model, _row_of(model, 1)) == SAMPLE_STATUS_FLAGGED_BG.lower()

    # Measured is automatic: rotation_count 1 counts, the 0s of 1-3 don't.
    assert _status(model, _row_of(model, 4)) == SAMPLE_STATUS_MEASURED_BG.lower()
    assert _status(model, _row_of(model, 2)) is None

    # Mounted always wins.
    model.updateCurrentSample(current_puck="P1", current_sample=1)
    assert _status(model, _row_of(model, 1)) == SAMPLE_ROW_QUEUED_BG.lower()


def test_status_filter_selects_rows(status_model):
    model = status_model
    model.set_queued_ids({1, 2})
    model.set_flagged(3, True)

    model.set_status_filter("queued")
    assert {model.get_id(r).db_id for r in range(model.rowCount())} == {1, 2}
    model.set_status_filter("flagged")
    assert {model.get_id(r).db_id for r in range(model.rowCount())} == {3}
    model.set_status_filter("measured")
    assert {model.get_id(r).db_id for r in range(model.rowCount())} == {4}
    # Unmeasured is the complement: everything still to be done.
    model.set_status_filter("unmeasured")
    assert {model.get_id(r).db_id for r in range(model.rowCount())} == {1, 2, 3}
    model.set_status_filter(None)
    assert model.rowCount() == 4


def test_queued_view_shows_queue_order(status_model):
    model = status_model
    # Queue order deliberately different from location/db order.
    model.set_queued_ids([3, 1, 2])
    model.set_status_filter("queued")
    assert [model.get_id(r).db_id for r in range(model.rowCount())] == [3, 1, 2]
    # Header sorts must not reorder the queue view — row 1 runs next.
    model.sort(1, Qt.SortOrder.AscendingOrder)
    assert [model.get_id(r).db_id for r in range(model.rowCount())] == [3, 1, 2]
    # Other views keep the normal header sort.
    model.set_status_filter(None)
    assert [model.get_id(r).db_id for r in range(model.rowCount())] != [3, 1, 2]


def test_status_tints_are_context_dependent(status_model):
    from aare.gui.styles import SAMPLE_STATUS_FLAGGED_BG, SAMPLE_STATUS_QUEUED_BG

    model = status_model
    model.set_queued_ids({1, 2})
    model.set_flagged(1, True)

    # Queued view: own tint suppressed, only the also-flagged mark shows.
    model.set_status_filter("queued")
    assert _status(model, _row_of(model, 1)) == SAMPLE_STATUS_FLAGGED_BG.lower()
    assert _status(model, _row_of(model, 2)) is None

    # Flagged view: a re-queued sample wears the queued mark.
    model.set_status_filter("flagged")
    assert _status(model, _row_of(model, 1)) == SAMPLE_STATUS_QUEUED_BG.lower()


def test_status_sets_refilter_while_chip_active(status_model):
    model = status_model
    model.set_status_filter("queued")
    assert model.rowCount() == 0
    model.set_queued_ids({2})
    assert {model.get_id(r).db_id for r in range(model.rowCount())} == {2}


def test_status_column_is_display_only(status_model):
    model = status_model
    assert model.data(model.index(0, 0), Qt.ItemDataRole.DisplayRole) == 1
    before = [model.get_id(r).db_id for r in range(model.rowCount())]
    model.sort(0, Qt.SortOrder.DescendingOrder)  # no-op on the "#" column
    assert [model.get_id(r).db_id for r in range(model.rowCount())] == before


def test_mime_data_round_trips_for_chip_drops(status_model):
    from aarecommon.models.models import SampleShortInfoList

    model = status_model
    payload = model.mimeData([model.index(0, 1), model.index(1, 1)])
    samples = SampleShortInfoList.model_validate_json(payload.text())
    assert len(samples.s) == 2
    assert samples.s[0].db_id == model.get_id(0).db_id