from __future__ import annotations

import contextlib
import logging
from typing import TYPE_CHECKING, Any, Literal, overload

from prettyqt import constants, core
from prettyqt.utils import bidict, helpers, listdelegators, modelhelpers


if TYPE_CHECKING:
    from collections.abc import Iterator, Sequence


logger = logging.getLogger(__name__)

CheckIndexOptionStr = Literal[
    "none", "index_is_valid", "do_not_use_parent", "parent_is_invalid"
]

CHECK_INDEX_OPTIONS: bidict[
    CheckIndexOptionStr, core.QAbstractItemModel.CheckIndexOption
] = bidict(
    none=core.QAbstractItemModel.CheckIndexOption.NoOption,
    index_is_valid=core.QAbstractItemModel.CheckIndexOption.IndexIsValid,
    do_not_use_parent=core.QAbstractItemModel.CheckIndexOption.DoNotUseParent,
    parent_is_invalid=core.QAbstractItemModel.CheckIndexOption.ParentIsInvalid,
)

LayoutChangeHintStr = Literal["none", "vertical_sort", "horizontal_sort"]

LAYOUT_CHANGE_HINT: bidict[
    LayoutChangeHintStr, core.QAbstractItemModel.LayoutChangeHint
] = bidict(
    none=core.QAbstractItemModel.LayoutChangeHint.NoLayoutChangeHint,
    vertical_sort=core.QAbstractItemModel.LayoutChangeHint.VerticalSortHint,
    horizontal_sort=core.QAbstractItemModel.LayoutChangeHint.HorizontalSortHint,
)


class AbstractItemModelMixin(core.ObjectMixin):
    IS_RECURSIVE = False  # for docs
    DELEGATE_DEFAULT = None
    DEFAULT_FLAGS = (
        constants.DRAG_ENABLED | constants.IS_ENABLED | constants.IS_SELECTABLE
    )

    def __repr__(self):
        return f"{type(self).__name__}: ({self.rowCount()}, {self.columnCount()})"

    def __len__(self) -> int:
        return self.rowCount()

    def __add__(self, other: core.QAbstractItemModel) -> core.ConcatenateTablesProxyModel:
        proxy = core.ConcatenateTablesProxyModel()
        proxy.addSourceModel(self)
        proxy.addSourceModel(other)
        return proxy

    @overload
    def __getitem__(self, index: tuple[int, int] | int) -> core.ModelIndex: ...

    @overload
    def __getitem__(
        self, index: tuple[slice, int] | tuple[int, slice] | tuple[slice, slice]
    ) -> listdelegators.ListDelegator[core.ModelIndex]: ...

    def __getitem__(
        self, index: tuple[int | slice, int | slice]
    ) -> core.ModelIndex | listdelegators.ListDelegator[core.ModelIndex]:
        # TODO: do proxies need mapToSource here?
        rowcount = self.rowCount()
        colcount = self.columnCount()
        match index:
            case int() as row, int() as col:
                if row >= rowcount or col >= rowcount:
                    raise IndexError(index)
                return self.index(row, col)
            case (row, col):
                indexes = [
                    self.index(i, j)
                    for i, j in helpers.iter_positions(row, col, rowcount, colcount)
                ]
                return listdelegators.ListDelegator(indexes)
            case int() as row:
                if row >= rowcount:
                    raise IndexError(index)
                # this here breaks PySide6 IPython test...
                return self.index(row, 0)
            case _:
                raise TypeError(index)

    def set_data(
        self,
        index: tuple[int | slice, int | slice] | core.ModelIndex,
        value: Any,
        role=constants.EDIT_ROLE,
    ):
        match index:
            case core.ModelIndex():
                self.setData(index, value, role)
            case (row, col):
                rowcount = self.rowCount()
                colcount = self.columnCount()
                for i, j in helpers.iter_positions(row, col, rowcount, colcount):
                    self.setData(self.index(i, j), value, role)
            case _:
                raise TypeError(index)

    def check_index(
        self,
        index: core.ModelIndex,
        *,
        index_is_valid: bool = False,
        do_not_use_parent: bool = False,
        parent_is_invalid: bool = False,
    ) -> bool:
        flag = core.QAbstractItemModel.CheckIndexOption.NoOption
        if index_is_valid:
            flag |= core.QAbstractItemModel.CheckIndexOption.IndexIsValid
        if do_not_use_parent:
            flag |= core.QAbstractItemModel.CheckIndexOption.DoNotUseParent
        if parent_is_invalid:
            flag |= core.QAbstractItemModel.CheckIndexOption.ParentIsInvalid
        check_flag = core.QAbstractItemModel.CheckIndexOption(0) | flag
        return self.checkIndex(index, check_flag)

    @contextlib.contextmanager
    def change_layout(self):
        """Context manager to change the layout.

        wraps calls with correct signals
        emitted at beginning: layoutAboutToBeChanged
        emitted at end: layoutChanged

        """
        self.layoutAboutToBeChanged.emit()
        yield None
        self.layoutChanged.emit()

    @contextlib.contextmanager
    def reset_model(self):
        """Context manager to reset the model.

        wraps calls with correct signals
        emitted at beginning: beginResetModel
        emitted at end: endResetModel

        """
        self.beginResetModel()
        yield None
        self.endResetModel()

    def get_column_type(
        self,
        column: int,
        *,
        rows_to_check: int = 10,
        role=constants.DISPLAY_ROLE,
    ) -> type | None:
        """Guess column data type by checking values of first rows with given role."""
        to_check = min(rows_to_check, self.rowCount())
        if to_check == 0:
            return None
        # cant combine these or make them a generator, so we do two list comps.
        indexes = [self.index(row, column) for row in range(to_check)]
        values = [self.data(i, role=role) for i in indexes]
        if all(isinstance(i, bool) for i in values):
            return bool
        if all(isinstance(i, str) for i in values):
            return str
        if all(isinstance(i, int) for i in values):
            return int
        if all(isinstance(i, float) for i in values):
            return float
        check_values = [self.data(i, role=constants.CHECKSTATE_ROLE) for i in indexes]
        return bool if None not in check_values else None

    def update_row(self, row: int):
        start_index = self.index(row, 0)
        end_index = self.index(row, self.columnCount() - 1)
        self.dataChanged.emit(start_index, end_index)

    def update_all(self):
        top_left = self.index(0, 0)
        bottom_right = self.index(self.rowCount() - 1, self.columnCount() - 1)
        self.dataChanged.emit(top_left, bottom_right)

    @contextlib.contextmanager
    def remove_row(self, row: int, parent: core.ModelIndex | None = None):
        parent = core.ModelIndex() if parent is None else parent
        self.beginRemoveRows(parent, row, row)
        yield None
        self.endRemoveRows()

    @contextlib.contextmanager
    def remove_rows(
        self,
        first: int | None = None,
        last: int | None = None,
        parent: core.ModelIndex | None = None,
    ):
        parent = core.ModelIndex() if parent is None else parent
        first = first or 0
        last = last if last is not None else self.rowCount()
        self.beginRemoveRows(parent, first, last)
        yield None
        self.endRemoveRows()

    @contextlib.contextmanager
    def remove_columns(
        self,
        first: int | None = None,
        last: int | None = None,
        parent: core.ModelIndex | None = None,
    ):
        parent = core.ModelIndex() if parent is None else parent
        first = first or 0
        last = last if last is not None else self.rowCount()
        self.beginRemoveColumns(parent, first, last)
        yield None
        self.endRemoveColumns()

    @contextlib.contextmanager
    def insert_row(self, row: int, parent: core.ModelIndex | None = None):
        parent = core.ModelIndex() if parent is None else parent
        self.beginInsertRows(parent, row, row)
        yield None
        self.endInsertRows()

    @contextlib.contextmanager
    def insert_rows(
        self,
        first: int | None = None,
        last: int | None = None,
        parent: core.ModelIndex | None = None,
    ):
        parent = core.ModelIndex() if parent is None else parent
        first = first or 0
        last = last if last is not None else self.rowCount()
        self.beginInsertRows(parent, first, last)
        yield None
        self.endInsertRows()

    @contextlib.contextmanager
    def append_rows(self, num_rows: int, parent: core.ModelIndex | None = None):
        parent = core.ModelIndex() if parent is None else parent
        row_count = self.rowCount()
        self.beginInsertRows(parent, row_count, row_count + num_rows - 1)
        yield None
        self.endInsertRows()

    @contextlib.contextmanager
    def insert_columns(
        self,
        first: int | None = None,
        last: int | None = None,
        parent: core.ModelIndex | None = None,
    ):
        parent = core.ModelIndex() if parent is None else parent
        first = first or 0
        last = last if last is not None else self.rowCount()
        self.beginInsertColumns(parent, first, last)
        yield None
        self.endInsertColumns()

    def force_reset(self):
        self.beginResetModel()
        self.endResetModel()

    def force_layoutchange(self):
        self.layoutAboutToBeChanged.emit()
        self.layoutChanged.emit()

    def get_role_names(self) -> bidict[int, str]:
        return bidict({i: v.data().decode() for i, v in self.roleNames().items()})

    def get_breadcrumbs_path(
        self,
        index: core.ModelIndex,
        role: constants.ItemDataRole = constants.DISPLAY_ROLE,
    ) -> list[str]:
        """Get the path for the given index.

        Returns a list containing data of all indexes up to the root for given role.
        """
        pieces = [index.data(role)]
        while (index := index.parent()).isValid():
            pieces.insert(0, index.data(role))
        return pieces

    def iter_tree(
        self,
        parent_index: core.ModelIndex | None = None,
        *,
        depth: int | None = None,
        fetch_more: bool = False,
    ) -> Iterator[core.ModelIndex]:
        """Iter through all indexes of the model tree.

        Arguments:
            parent_index: parent index
            depth: maximum iteration depth
            fetch_more: call fetchMore for all indexes until canFetchMore returns False
        """
        if parent_index is None:
            # TODO: does this always equal AbstractItemView.rootIndex()?
            # parent_index = self.index(0, 0)
            parent_index = core.ModelIndex()
        if parent_index.isValid():
            yield parent_index
            if fetch_more:
                while self.canFetchMore(parent_index):
                    self.fetchMore(parent_index)
        if depth is not None and (depth := depth - 1) < 0:
            return
        for i in range(self.rowCount(parent_index)):
            child_index = self.index(i, 0, parent_index)
            yield from self.iter_tree(child_index, depth=depth, fetch_more=fetch_more)

    def search_tree(
        self,
        value: Any,
        *,
        role: constants.ItemDataRole = constants.DISPLAY_ROLE,
        parent_index: core.ModelIndex | None = None,
        max_results: int | None = None,
        depth: int | None = None,
    ) -> listdelegators.ListDelegator[core.ModelIndex]:
        """Search the tree for indexes with a given value in given role.

        Compared to QAbstractItemModel.match, this method allows to set a maximum
        search depth and passing several values to search for as a list.

        Arguments:
            value: Item or list of items to search for.
            role: Index role to search in.
            parent_index: start index for searching. If None, whole tree is searched.
            max_results: stop searching after x amount of hits. 'None' means no limit.
            depth: search depth. Search depth. 'None' means no limit.
        """
        results = []
        # This makes it impossible to search for lists. I think thats fine.
        if not isinstance(value, list):
            value = [value]
        for idx in self.iter_tree(parent_index, depth=depth):
            if self.data(idx, role) in value:
                results.append(idx)
                if len(results) == max_results:
                    break
        return listdelegators.ListDelegator(results)

    def get_child_indexes(
        self, index: core.ModelIndex
    ) -> listdelegators.ListDelegator[core.ModelIndex]:
        """Get all child indexes for given index (first column only).

        To get indexes recursively, use iter_tree.

        Arguments:
            index: ModelIndex to get children from
        """
        indexes = [self.index(i, 0, index) for i in range(self.rowCount(index))]
        return listdelegators.ListDelegator(indexes)

    def get_index_key(
        self,
        index: core.ModelIndex,
        *,
        include_column: bool = False,
        parent_index: core.ModelIndex | None = None,
    ) -> tuple[tuple[int, int], ...]:
        """Return a key tuple for given ModelIndex.

        The key tuple consists either of row integers or (row, column) indices
        describing the index position from top to bottom.

        Arguments:
            index: ModelIndex to get a key for
            include_column: whether to include the column in the index key.
            parent_index: Get key up to given ModelIndex. By default, get key up to root.
        """
        key_path = []
        parent = index
        while parent.isValid() and parent != parent_index:
            key = (parent.row(), parent.column()) if include_column else parent.row()
            key_path.append(key)
            parent = parent.parent()
        return tuple(reversed(key_path))

    def index_from_key(
        self,
        key_path: Sequence[tuple[int, int] | int],
        *,
        parent_index: core.ModelIndex | None = None,
    ) -> core.ModelIndex:
        """Return a  ModelIndex for the given key path.

        Arguments:
            key_path: Key path to get an index for.
                      Should be a sequence of either (row, column)  or row indices
            parent_index: ModelIndex to start indexing from. Defaults to root index.
        """
        return modelhelpers.index_from_key(self, key_path, parent_index=parent_index)

    @staticmethod
    def to_checkstate(value: bool):
        return constants.CheckState.Checked if value else constants.CheckState.Unchecked

    def get_header_index_for_value(
        self,
        label,
        orientation: (
            constants.Orientation | constants.OrientationStr
        ) = constants.HORIZONTAL,
        *,
        role: constants.ItemDataRole = constants.DISPLAY_ROLE,
    ) -> int | None:
        o = constants.ORIENTATION.get_enum_value(orientation)
        count = self.columnCount() if o == constants.HORIZONTAL else self.rowCount()
        return next(
            (i for i in range(count) if self.headerData(i, o, role) == label), None
        )

    def get_table_data(
        self,
        *,
        include_index: bool = False,
        role: constants.ItemDataRole = constants.DISPLAY_ROLE,
        x_range: slice | int | None = None,
        y_range: slice | int | None = None,
        parent_index: core.ModelIndex | None = None,
        use_checkstate_role: bool = False,
    ):
        parent_index = parent_index or core.ModelIndex()
        # if self.canFetchMore(parent_index):
        #     self.fetchMore(parent_index)
        match x_range:
            case None:
                colrange = range(self.columnCount(parent_index))
            case slice():
                stop = x_range.stop or self.columnCount(parent_index)
                colrange = range(x_range.start or 0, stop, x_range.step or 1)
            case int():
                colrange = range(x_range, x_range + 1)
            case _:
                raise TypeError(x_range)

        match y_range:
            case None:
                rowrange = range(self.rowCount(parent_index))
            case slice():
                stop = y_range.stop or self.rowCount(parent_index)
                rowrange = range(y_range.start or 0, stop, y_range.step or 1)
            case int():
                rowrange = range(y_range, y_range + 1)
            case _:
                raise TypeError(y_range)
        data = [
            [self.index(i, j, parent_index).data(role) for j in colrange]
            for i in rowrange
        ]
        if use_checkstate_role:
            check_data, _, __ = self.get_table_data(
                role=constants.CHECKSTATE_ROLE,
                x_range=x_range,
                y_range=y_range,
                parent_index=parent_index,
            )
            for i, row in enumerate(data):
                for j, _column in enumerate(row):
                    if check_data[i][j] is not None:
                        data[i][j] = bool(check_data[i][j])
        h_header = [self.headerData(i, constants.HORIZONTAL) for i in colrange]
        v_header = (
            [self.headerData(i, constants.VERTICAL) for i in rowrange]
            if include_index
            else None
        )
        return data, h_header, v_header

    def to_dataframe(self, **kwargs):
        import pandas as pd

        data, h_header, v_header = self.get_table_data(**kwargs)
        return pd.DataFrame(data=data, columns=h_header, index=v_header)


class AbstractItemModel(AbstractItemModelMixin, core.QAbstractItemModel):
    """The abstract interface for item model classes."""

    # @abc.abstractmethod
    # def index(self, *args, **kwargs):
    #     return NotImplemented

    # # this one is only abstract for a specific signature.
    # # @abc.abstractmethod
    # # def parent(self, *args, **kwargs):
    # #     return NotImplemented

    # @abc.abstractmethod
    # def rowCount(self, *args, **kwargs):
    #     return NotImplemented

    # @abc.abstractmethod
    # def columnCount(self, *args, **kwargs):
    #     return NotImplemented

    # @abc.abstractmethod
    # def data(self, *args, **kwargs):
    #     return NotImp


if __name__ == "__main__":

    class Test(AbstractItemModel):
        pass
