from PySide6.QtCore import QAbstractTableModel, Qt
from PySide6.QtGui import QBrush, QColor

from aare.common.models import SampleShortInfo, SampleShortInfoList


def get_entry(sample: SampleShortInfo, column: int, *, show_user: bool = False):
    if show_user:
        if column == 0:
            return sample.user
        elif column == 1:
            return sample.dewar_name
        elif column == 2:
            return sample.puck_name
        elif column == 3:
            return sample.sample_name
    else:
        if column == 0:
            return sample.dewar_name
        elif column == 1:
            return sample.puck_name
        elif column == 2:
            return sample.sample_name
    return ""


class SampleQueueSpreadsheet(QAbstractTableModel):
    def __init__(
        self,
        parent=None,
        samples: list[SampleShortInfo] | None = None,
        show_user: bool = False,
    ):
        super().__init__(parent)
        if samples is None:
            samples = []
        self.__running = False
        self._show_user = bool(show_user)
        self.samples: list[SampleShortInfo] = samples
        self.header = (
            ["User", "Dewar", "Puck", "Sample"]
            if self._show_user
            else ["Dewar", "Puck", "Sample"]
        )

    def flags(self, index):
        default_flags = super().flags(index)
        return Qt.ItemFlag.ItemIsDropEnabled | default_flags

    def rowCount(self, parent=None):
        return len(self.samples)

    def columnCount(self, parent=None):
        return len(self.header)

    def data(self, index, role=None):
        if role == Qt.ItemDataRole.DisplayRole:
            return get_entry(self.samples[index.row()], index.column(), show_user=self._show_user)
        elif role == Qt.ItemDataRole.TextAlignmentRole:
            return Qt.AlignmentFlag.AlignCenter
        elif role == Qt.ItemDataRole.BackgroundRole:
            if index.row() == 0:
                if self.__running:
                    return QBrush(QColor(255, 102, 0))
                else:
                    return QBrush(QColor(114, 159, 207))
            return QBrush(QColor(255, 255, 255))
        return None

    def headerData(self, section, orientation, role=None):
        if role == Qt.ItemDataRole.DisplayRole:
            if orientation == Qt.Orientation.Horizontal:
                return self.header[section] if self.header else f"Column {section + 1}"
            if orientation == Qt.Orientation.Vertical:
                return str(section + 1)
        return None

    def updateData(self, samples: list[SampleShortInfo],):
        self.beginResetModel()
        self.samples = list(samples)
        self.endResetModel()

    def mimeTypes(self):
        return ['text/plain']

    def canDropMimeData(self, data, action, row, column, parent):
        if data.hasText():
            return True
        return False

    def dropMimeData(self, data, action, row, column, parent):
        if not self.canDropMimeData(data, action, row, column, parent):
            return False

        if parent.isValid():
            row = parent.row()

        try:
            l = SampleShortInfoList.model_validate_json(data.text())
            self.beginResetModel()
            for sample in l.s:
                updated_row = row
                updated_samples = []
                for i in range(len(self.samples)):
                    if sample != self.samples[i]:
                        updated_samples.append(self.samples[i])
                    else:
                        if 0 < row < i:
                            updated_row -= 1
                if updated_row == -1:
                    updated_samples.append(sample)
                else:
                    updated_samples.insert(updated_row, sample)
                self.samples = updated_samples
            self.endResetModel()
        except Exception as e:
            print(e)
            return False
        return True

    def set_running(self, running: bool):
        self.__running = running
        self.layoutChanged.emit()

    def remove_sample(self, db_id: int):
        self.beginResetModel()
        self.samples = list(filter(lambda x: x.db_id != db_id, self.samples))
        self.endResetModel()

    def clearSamples(self):
        self.beginResetModel()
        self.samples = []
        self.endResetModel()

    def annotate_sample(self, db_id: int, comment: str) -> None:
        self.beginResetModel()
        updated_samples: list[SampleShortInfo] = []
        for sample in self.samples:
            if sample.db_id == db_id:
                updated_samples.append(sample.model_copy(update={"comment": comment}))
            else:
                updated_samples.append(sample)
        self.samples = updated_samples
        self.endResetModel()

    def to_state(self) -> dict:
        return {
            "samples": [s.model_dump() for s in self.samples],
        }

    def from_state(self, state: dict):
        self.beginResetModel()

        self.samples = [
            SampleShortInfo.model_validate(d)
            for d in state.get("samples", [])
        ]

        self.endResetModel()