import typing
from functools import partial

import pandas as pd
from PyQt6.QtGui import QPixmap, QAction
from PyQt6.QtWidgets import QMainWindow, QTableView, QHBoxLayout, QLabel, \
    QMenu, QHeaderView, QLineEdit
from PyQt6 import uic, QtGui, QtCore, QtWidgets
from PyQt6.QtCore import Qt, QSortFilterProxyModel

from banking_breakdown.ui.main_window import Ui_MainWindow


class PandasModel(QtCore.QAbstractTableModel):
    def __init__(self, df: pd.DataFrame, parent=None):
        QtCore.QAbstractTableModel.__init__(self, parent)
        self._data = df

        self._horizontalHeaders = [''] * len(df.columns)

        for i, column in enumerate(df.columns):
            self.setHeaderData(i, Qt.Orientation.Horizontal, column)

    def setHeaderData(self, section, orientation, data,
                      role=Qt.ItemDataRole.EditRole):

        if ((orientation == Qt.Orientation.Horizontal)
                and ((role == Qt.ItemDataRole.DisplayRole)
                     or (role == Qt.ItemDataRole.EditRole))):

            self._horizontalHeaders[section] = data
            return True
        else:
            return super().setHeaderData(section, orientation, data, role)

    def headerData(self, section, orientation,
                   role=Qt.ItemDataRole.DisplayRole):

        if (orientation == Qt.Orientation.Horizontal
                and role == Qt.ItemDataRole.DisplayRole):

            return self._horizontalHeaders[section]
        else:
            return super().headerData(section, orientation, role)

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

    def columnCount(self, parent=None):
        return self._data.columns.size

    def data(self, index, role=Qt.ItemDataRole.DisplayRole):
        if not index.isValid():
            return QtCore.QVariant()
        if role != Qt.ItemDataRole.DisplayRole:
            return QtCore.QVariant()

        item = self._data.iloc[index.row()].iloc[index.column()]

        if type(item) is pd.Timestamp:
            return QtCore.QVariant(item.strftime('%Y-%m-%d'))
        else:
            return QtCore.QVariant(str(item))


class GeneratedWindowWrapper(QMainWindow):
    def __init__(self):
        super(GeneratedWindowWrapper, self).__init__()

        self._window = Ui_MainWindow()
        self._window.setupUi(self)

        # Populate categories

        self._category_model = QtGui.QStandardItemModel()
        self._window.categoryListView.setModel(self._category_model)

        # Onclick handlers

        self._window.createCategoryButton.clicked.connect(
            self._handle_create_click)

        header = self._window.transactionTableView.horizontalHeader()
        header.setContextMenuPolicy(Qt.ContextMenuPolicy.CustomContextMenu)
        header.customContextMenuRequested.connect(self._header_right_clicked)

    def _header_right_clicked(self, pos):
        context = QMenu(self)
        context.addAction(QAction("test 1", self))
        context.addAction(QAction("test 2", self))
        context.addAction(QAction("test 3", self))
        context.exec(self.sender().mapToGlobal(pos))

    def add_categories(self, categories: typing.Sequence[str]):
        for category in categories:
            item = QtGui.QStandardItem(category)
            self._category_model.appendRow(item)

    def set_statement_data(self, df: pd.DataFrame):
        model = PandasModel(df)
        proxyModel = QSortFilterProxyModel(self)
        proxyModel.setSourceModel(model)
        self._window.transactionTableView.setModel(proxyModel)

        self._window.transactionTableView.resizeColumnsToContents()

    def _handle_create_click(self):
        self.add_categories(['asdf'])

    def _handle_delete_click(self):
        pass

    def _handle_apply_click(self):
        pass

    def add_warning_text(self, text: str):
        layout = QHBoxLayout()

        warningIcon = QLabel()
        pixmap = QPixmap('res/warning.png')
        warningIcon.setPixmap(pixmap)

        label = QLabel(text)
        label.setWordWrap(True)

        layout.addWidget(warningIcon)
        layout.addWidget(label)

        layout.setStretch(0, 0)
        layout.setStretch(1, 1)

        self._window.warningLayout.addLayout(layout)