from typing import Any, List, Union

from qgis.PyQt import QtCore
from qgis.PyQt.QtCore import QModelIndex, QObject, QSortFilterProxyModel, Qt, QVariant

from openlog.datamodel.assay.generic_assay import AssayDefinition
from openlog.datamodel.connection.openlog_connection import OpenLogConnection
from openlog.gui.assay_visualization.config.assay_column_visualization_config import (
    AssayColumnVisualizationConfig,
)
from openlog.gui.assay_visualization.config.assay_visualization_config import (
    AssayVisualizationConfig,
    AssayVisualizationConfigList,
)
from openlog.gui.assay_visualization.discrete.config import (
    DiscreteAssayColumnVisualizationConfig,
    DiscreteCategoricalAssayColumnVisualizationConfig,
)
from openlog.gui.assay_visualization.extended.config import (
    ExtendedCategoricalAssayColumnVisualizationConfig,
    ExtendedNumericalAssayColumnVisualizationConfig,
)
from openlog.gui.assay_visualization.stacked.stacked_config import (
    StackedConfiguration,
    StackedConfigurationList,
)
from openlog.gui.check_state_model import CheckStateModel


class AssayVisualizationConfigProxyModel(QSortFilterProxyModel):
    def __init__(self, parent: QObject = None):
        super().__init__(parent)

    def lessThan(self, left: QModelIndex, right: QModelIndex) -> bool:
        if left.parent().isValid():
            return super().lessThan(left, right)
        else:
            return False


class AssayVisualizationConfigTreeModel(CheckStateModel):

    CONFIG_ROLE = QtCore.Qt.UserRole
    DISPLAY_NAME_ROLE = QtCore.Qt.UserRole + 1
    HOLE_ID_ROLE = QtCore.Qt.UserRole + 2

    def __init__(
        self,
        assay_visualization_config_list: AssayVisualizationConfigList,
        stacked_config_list: StackedConfigurationList,
        openlog_connection: OpenLogConnection,
        parent: QObject = None,
    ):
        """
        Qt TreeModel implementation for AssayVisualizationConfigList and StackedConfigurationList display and edition.
        Configuration is displayed with 3 parent item :
        - collar display : _collar_root_index
        - assay display : _assay_root_index
        - stacked display : _stacked_root_index

        The configuration are synchronized by setData override:
         - all configuration are searched and we apply new check state for each configuration indexes

        To display collar and assay without associated configuration, use :
        - add_assay / remove_assay
        - add_collar / remove_collar


        Args:
            assay_visualization_config_list: AssayVisualizationConfigList to be displayed
            openlog_connection: OpenLogConnection used to get collar display name
            parent: parent object
        """
        super().__init__(0, 1, parent)
        self.assay_visualization_config_list = assay_visualization_config_list
        self.stacked_config_list = stacked_config_list
        self.openlog_connection = openlog_connection
        self.setHorizontalHeaderLabels([self.tr("Name")])
        self.refresh()

    @property
    def _collar_root_index(self) -> QModelIndex:
        res = self._get_root_child_from_name(self.tr("Collars"))
        if not res.isValid():
            res = self._insert_root_item(self.tr("Collars"))
        return res

    @property
    def _assay_root_index(self) -> QModelIndex:
        res = self._get_root_child_from_name(self.tr("Downhole data"))
        if not res.isValid():
            res = self._insert_root_item(self.tr("Downhole data"))
        return res

    @property
    def _stacked_root_index(self) -> QModelIndex:
        res = self._get_root_child_from_name(self.tr("Stacked"))
        if not res.isValid():
            res = self._insert_root_item(self.tr("Stacked"))
        return res

    def refresh(self):
        """
        Update tree model with current AssayVisualizationConfigList.

        """
        self._clear_unused_collar(self.assay_visualization_config_list.collar_list())
        self._clear_unused_assay(
            self.assay_visualization_config_list.assay_display_name_list()
        )

        for assay_config in self.assay_visualization_config_list.list:
            self._add_assay_visualization_config(assay_config)

        for stacked_config in self.stacked_config_list.list:
            self._add_stacked_config(stacked_config)

    def data(
        self, index: QtCore.QModelIndex, role: int = QtCore.Qt.DisplayRole
    ) -> QVariant:
        """
        Override QStandardItemModel data() for collar name display instead of collar id

        Args:
            index: QModelIndex
            role: Qt role

        Returns: QVariant

        """
        result = super().data(index, role)
        if (
            self.openlog_connection
            and role == QtCore.Qt.DisplayRole
            and index.column() == 0
            and result
            and super().data(index, self.DISPLAY_NAME_ROLE)
        ):
            result = super().data(index, self.DISPLAY_NAME_ROLE)

        return result

    def _clear_unused_collar(self, available_collar: List[str]) -> None:
        """
        Remove unused collar from _assay_root_index

        Args:
            available_collar: available collars
        """
        for i in range(0, self.rowCount(self._assay_root_index)):
            self._remove_item_not_available(
                self._assay_root_index.child(i, 0), available_collar
            )

    def _clear_unused_assay(self, available_assay: List[str]) -> None:
        """
        Remove unused assay from _collar_root_index

        Args:
            available_assay: available assay
        """
        for i in range(0, self.rowCount(self._collar_root_index)):
            self._remove_item_not_available(
                self._collar_root_index.child(i, 0), available_assay
            )

    def add_assay(self, assay_def: AssayDefinition) -> None:
        """
        Add assay item in _assay_root_index (no add if assay already available)

        Args:
            assay_def: (AssayDefinition) assay definition
        """
        index, _ = self._get_or_insert_item(self._assay_root_index, assay_def.variable)
        super().setData(index, assay_def.display_name, self.DISPLAY_NAME_ROLE)

    def remove_assay(self, assay_name: str) -> None:
        """
        Remove assay item in _assay_root_index (no error if assay not available)

        Args:
            assay_name: (str) assay name
        """
        self._remove_item_by_name(self._assay_root_index, assay_name)

    def add_collar(self, collar: str) -> None:
        """
        Add collar item in _collar_root_index (no add if collar already available)

        Args:
            collar: (str) collar name
        """
        collar_display_name = (
            self.openlog_connection.get_read_iface().get_collar_display_name(collar)
        )
        index, _ = self._get_or_insert_item(self._collar_root_index, collar)
        super().setData(index, collar_display_name, self.DISPLAY_NAME_ROLE)
        super().setData(index, collar, self.HOLE_ID_ROLE)

    def remove_collar(self, hole_id: str) -> None:
        """
        Remove collar item in _collar_root_index (no error if collar not available)

        Args:
            hole_id: (str) collar name
        """
        self._remove_item_by_name(self._collar_root_index, hole_id)

    def get_available_assay(self, only_visible: bool = False) -> List[str]:
        """
        Return assay available in tree model.

        Args:
            only_visible: True to returns only visible assay (default False)

        Returns: list of assay name

        """
        return self._get_available_items(self._assay_root_index, only_visible)

    def get_available_collar(self, only_visible: bool = False) -> List[str]:
        """
        Return collar available in tree model.

        Args:
            only_visible: True to returns only visible collar (default False)

        Returns: list of collar name

        """
        return self._get_available_items(self._collar_root_index, only_visible)

    def is_assay(self, index: QModelIndex) -> bool:
        """
        Check if an index is used to display an assay

        Args:
            index: QModelIndex

        Returns: True if index is used to display an assay, False otherwise

        """
        return index.parent() == self._assay_root_index

    def is_collar(self, index: QModelIndex) -> bool:
        """
        Check if an index is used to display a collar

        Args:
            index: QModelIndex

        Returns: True if index is used to display an assay, False otherwise

        """
        return index.parent() == self._collar_root_index or self.is_assay(
            index.parent()
        )

    def is_assay_column(self, index: QModelIndex) -> bool:
        """
        Check if an index is used to display an assay column (when multiple column within an assay)

        Args:
            index: QModelIndex

        Returns: True if index is used to display an assay column, False otherwise

        """
        return (
            index.parent().parent().parent() == self._assay_root_index
            or index.parent().parent().parent() == self._collar_root_index
        )

    def remove_stacked(self, config: StackedConfiguration) -> None:
        for i in range(0, self.rowCount(self._stacked_root_index)):
            if (
                self.data(self.index(i, 0, self._stacked_root_index), self.CONFIG_ROLE)
                == config
            ):
                self.removeRow(i, self._stacked_root_index)
                break

    def is_stacked(self, index: QModelIndex) -> bool:
        """
        Check if an index is used to display a stacked

        Args:
            index: QModelIndex

        Returns: True if index is used to display an stacked, False otherwise

        """
        return index.parent() == self._stacked_root_index

    def is_splittable(self, index: QModelIndex) -> bool:
        """
        Check if an index is splittable for display in canvas

        Args:
            index: QModelIndex

        Returns: True if item is splittable, False otherwise

        """

        splittable_types = (
            DiscreteAssayColumnVisualizationConfig,
            DiscreteCategoricalAssayColumnVisualizationConfig,
            ExtendedCategoricalAssayColumnVisualizationConfig,
            ExtendedNumericalAssayColumnVisualizationConfig,
        )
        if self.is_assay(index):
            config = index.data(self.CONFIG_ROLE)
            column = list(config.column_config.values())[0]
            return isinstance(column, splittable_types)
        elif self.is_assay_column(index):
            config = index.data(self.CONFIG_ROLE)
            return isinstance(config, splittable_types)
        else:
            return False

    def flags(self, index: QtCore.QModelIndex) -> QtCore.Qt.ItemFlags:
        """
        Override QStandardItemModel flags.

        Args:
            index: QModelIndex

        Returns: index flags

        """
        # All item are enabled and selectable
        flags = Qt.ItemFlag.ItemIsEnabled | Qt.ItemFlag.ItemIsSelectable
        # All item should be checkable
        flags = flags | Qt.ItemIsUserCheckable | Qt.ItemIsAutoTristate
        return flags

    def setData(
        self, index: QtCore.QModelIndex, value: Any, role: int = Qt.DisplayRole
    ) -> bool:
        """
        Override QStandardItemModel setData for configuration synchronization.

        Args:
            index: QModelIndex
            value: new value
            role: Qt role

        Returns: True if data set, False otherwise

        """
        done = False
        res = False

        # Only sync QCheckStateRole for assay visibility
        if role == Qt.CheckStateRole:
            done = self._set_check_state_data(index, value, role)
        if not done:
            res = super().setData(index, value, role)

        return res

    def _set_check_state_data(
        self, index: QtCore.QModelIndex, value: Any, role: int = Qt.DisplayRole
    ) -> bool:
        """
        Define check state data

        Args:
            index: QModelIndex
            value: new value
            role: Qt role

        Returns: True if data set, False otherwise

        """
        done_ = False
        oldState = self.data(index, role)
        if value != oldState:
            config = self.data(index, self.CONFIG_ROLE)
            # Update configuration and plot visibility
            visible = value == Qt.Checked
            if isinstance(config, AssayVisualizationConfig):
                self._update_assay_visualization_config(visible, config)

                # Propagate change to all configuration indexes
                for all_config_index in self._get_config_indexes(config):
                    super().setData(all_config_index, value, role)

                done_ = True
            elif isinstance(config, AssayColumnVisualizationConfig):
                config.visibility_param.setValue(visible)
            elif isinstance(config, StackedConfiguration):
                config.is_visible = visible
                if config.plot:
                    config.plot.setVisible(visible)
                super().setData(index, value, role)
                done_ = True
        return done_

    @staticmethod
    def _update_assay_visualization_config(
        visible: bool, config: AssayVisualizationConfig
    ) -> None:
        """
        Update AssayVisualizationConfig for visibility and define plot widgets visibility

        Args:
            visible: (bool) define visibility
            config: (AssayVisualizationConfig) config to be updated
        """
        config.is_visible = visible
        if config.plot_stacked:
            config.plot_stacked.setVisible(visible and config.stacked_param.value())
        for col, col_config in config.column_config.items():
            if col_config.plot_widget:
                col_config.plot_widget.setVisible(
                    visible
                    and not config.stacked_param.value()
                    and col_config.visibility_param.value()
                )

    def _add_assay_visualization_config(
        self, assay_visualization_config: AssayVisualizationConfig
    ) -> None:
        """
        Add AssayVisualizationConfig items for collar index and assay index

        Args:
            assay_visualization_config: AssayVisualizationConfig
        """
        if (
            assay_visualization_config.hole_id
            and assay_visualization_config.assay_display_name
        ):
            collar_index, _, = self._get_or_insert_item(
                self._collar_root_index, assay_visualization_config.hole_id
            )
            index = self._insert_assay_visualization_config(
                collar_index,
                assay_visualization_config.assay_name,
                assay_visualization_config,
            )
            # set display name
            self.setData(
                index,
                assay_visualization_config.assay_display_name,
                self.DISPLAY_NAME_ROLE,
            )

            assay_index, _ = self._get_or_insert_item(
                self._assay_root_index, assay_visualization_config.assay_name
            )
            index = self._insert_assay_visualization_config(
                assay_index,
                assay_visualization_config.hole_id,
                assay_visualization_config,
            )

            # set hole display name
            self.setData(
                index,
                assay_visualization_config.hole_display_name,
                self.DISPLAY_NAME_ROLE,
            )
            # set hole_id
            self.setData(
                index,
                assay_visualization_config.hole_id,
                self.HOLE_ID_ROLE,
            )

        elif assay_visualization_config.assay_display_name:
            assay_index, _ = self._get_or_insert_item(
                self._assay_root_index, assay_visualization_config.assay_name
            )
            self.setData(assay_index, assay_visualization_config, self.CONFIG_ROLE)
            self.setData(
                assay_index,
                assay_visualization_config.assay_display_name,
                self.DISPLAY_NAME_ROLE,
            )

    def _add_stacked_config(self, stacked_config: StackedConfiguration) -> None:
        """
        Insert or update StackedConfiguration  with wanted name

        Args:
            stacked_config: StackedConfiguration
        """
        name = stacked_config.name
        index, _ = self._get_or_insert_item(self._stacked_root_index, name)
        self.setData(index, stacked_config, self.CONFIG_ROLE)
        super().setData(
            index,
            Qt.Checked if stacked_config.is_visible else Qt.Unchecked,
            Qt.CheckStateRole,
        )

    def _insert_assay_visualization_config(
        self,
        parent: QModelIndex,
        name: str,
        assay_visualization_config: AssayVisualizationConfig,
    ) -> QModelIndex:
        """
        Insert or update AssayVisualization for a parent index (collar or assay) with wanted name

        Args:
            parent: (QModelIndex) parent index
            name: (str) display name
            assay_visualization_config: AssayVisualizationConfig
        """
        index, _ = self._get_or_insert_item(parent, name)
        self.setData(index, assay_visualization_config, self.CONFIG_ROLE)

        super().setData(
            index,
            Qt.Checked if assay_visualization_config.is_visible else Qt.Unchecked,
            Qt.CheckStateRole,
        )

        # Insert columns configuration only if multiple columns are available
        if len(assay_visualization_config.column_config.items()) > 1:
            # Add index for all available columns
            for col, col_config in assay_visualization_config.column_config.items():
                self._insert_assay_column_visualization_config(index, col, col_config)

        return index

    def _insert_assay_column_visualization_config(
        self,
        parent: QModelIndex,
        name: str,
        assay_column_visualization_config: AssayColumnVisualizationConfig,
    ) -> None:
        """
        Insert or update AssayColumnVisualizationConfig for a parent index (collar or assay) with wanted name

        Args:
            parent: (QModelIndex) parent index
            name: (str) display name
            assay_column_visualization_config: AssayVisualizationConfig
        """
        index, inserted = self._get_or_insert_item(parent, name)
        self.setData(index, assay_column_visualization_config, self.CONFIG_ROLE)
        super().setData(
            index,
            Qt.Checked
            if assay_column_visualization_config.visibility_param.value()
            else Qt.Unchecked,
            Qt.CheckStateRole,
        )

        # Only connect visibility param signal once to avoid multiple connection
        if inserted:
            assay_column_visualization_config.visibility_param.sigValueChanged.connect(
                lambda: self._assay_column_visibility_param_changed(
                    assay_column_visualization_config
                )
            )

    def _assay_column_visibility_param_changed(
        self, assay_column_visualization_config: AssayColumnVisualizationConfig
    ) -> None:
        """
        Update model when visibility param is changed for a AssayColumnVisualizationConfig

        Args:
            assay_column_visualization_config: (AssayColumnVisualizationConfig) configuration
        """
        indexes = self._get_config_indexes(assay_column_visualization_config)
        visible = assay_column_visualization_config.visibility_param.value()
        for index in indexes:
            current_visibility = (
                True if self.data(index, Qt.CheckStateRole) == Qt.Checked else False
            )
            if visible != current_visibility:
                super().setData(
                    index,
                    Qt.Checked if visible else Qt.Unchecked,
                    Qt.CheckStateRole,
                )

    def _get_config_indexes(
        self, config: Union[AssayVisualizationConfig, AssayColumnVisualizationConfig]
    ) -> List[QModelIndex]:
        """
        Get AssayVisualizationConfig indexes

        Args:
            config: AssayVisualizationConfig

        Returns: list of indexes associated with config

        """
        res = self._get_config_indexes_from_parent(self._collar_root_index, config)
        res += self._get_config_indexes_from_parent(self._assay_root_index, config)
        return res

    def _get_config_indexes_from_parent(
        self,
        parent: QModelIndex,
        config: Union[AssayVisualizationConfig, AssayColumnVisualizationConfig],
    ) -> List[QModelIndex]:
        """
        Get AssayVisualizationConfig indexes from a parent index

        Args:
            parent: QModelIndex
            config: AssayVisualizationConfig

        Returns: list of indexes associated with config inside this parent

        """
        res = []
        for i in range(0, self.rowCount(parent)):
            index = parent.child(i, 0)
            if config == self.data(index, self.CONFIG_ROLE):
                res.append(index)
            res += self._get_config_indexes_from_parent(index, config)
        return res

    def _insert_root_item(self, name: str) -> QModelIndex:
        """
        Insert an item in root and returns created QModelIndex

        Args:
            name: item display name

        Returns: created item QModelIndex

        """
        # Add row in root
        row = self.rowCount()
        self.insertRow(row)
        index = self.index(row, 0)
        self.setData(index, name)

        # Insert column in index for child display
        self.insertColumn(0, index)
        return index

    def _get_root_child_from_name(self, name: str) -> QModelIndex:
        """
        Get root child from display name

        Args:
            name: child name

        Returns: invalid QModelIndex if child not found, child index otherwise

        """
        res = QModelIndex()
        for i in range(0, self.rowCount()):
            index = self.index(i, 0)
            if name == super().data(index, Qt.DisplayRole):
                res = index
        return res

    def _get_first_child_from_name(self, parent: QModelIndex, name: str) -> QModelIndex:
        """
        Get first parent child from display name

        Args:
            parent: parent QModelIndex
            name: child name

        Returns: invalid QModelIndex if child not found, child index otherwise

        """
        res = QModelIndex()
        for i in range(0, self.rowCount(parent)):
            index = parent.child(i, 0)
            if name == super().data(index, Qt.DisplayRole):
                res = index
        return res

    def _get_or_insert_item(
        self, parent: QModelIndex, name: str
    ) -> (QModelIndex, bool):
        """
        Add or insert item with name in parent

        Args:
            parent: parent QModelIndex
            name: display name

        Returns: item QModelIndex, bool True if item was inserted, False otherwise

        """
        inserted = False
        res = self._get_first_child_from_name(parent, name)
        if not res.isValid():
            # Add row in parent
            row = self.rowCount(parent)
            self.insertRow(row, parent)
            res = parent.child(row, 0)
            self.setData(res, name)
            self.setData(res, Qt.Checked, Qt.CheckStateRole)

            # Insert column in index for child display
            self.insertColumn(0, res)
            inserted = True
        return res, inserted

    def _remove_item_by_name(self, parent: QModelIndex, name: str) -> None:
        """
        Remove child from parent by name

        Args:
            parent: parent QModelIndex
            name: child name
        """
        index = self._get_first_child_from_name(parent, name)
        if index.isValid():
            self.removeRow(index.row(), parent)

    def _remove_item_not_available(
        self, parent: QModelIndex, available_names: List[str]
    ) -> None:
        """
        Remove item in parent with display name not in available list

        Args:
            parent: parent QModelIndex
            available_names: available name list
        """
        for i in reversed(range(0, self.rowCount(parent))):
            if super().data(parent.child(i, 0), Qt.DisplayRole) not in available_names:
                self.removeRow(i, parent)

    def _get_available_items(
        self, parent: QModelIndex, only_visible: bool = False
    ) -> List[str]:
        """
        Return list of child names from a parent.

        Args:
            parent: parent QModelIndex
            only_visible: True to returns only checked child (default False)

        Returns: list of child name

        """
        res = []
        for i in range(0, self.rowCount(parent)):
            index = parent.child(i, 0)
            name = super().data(index, Qt.DisplayRole)
            if only_visible:
                check_state = self.data(index, Qt.CheckStateRole)
                if check_state != Qt.Unchecked:
                    res.append(name)
            else:
                res.append(name)
        return res
