import sys
import typing

import pandas as pd
from PyQt5.QtWidgets import QMainWindow, QApplication, QPushButton, \
    QListView, QTableView
from PyQt5 import uic, QtGui, QtCore
import datetime


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, QtCore.Qt.Horizontal, column)

    def setHeaderData(self, section, orientation, data,
                      role=QtCore.Qt.EditRole):
        if orientation == QtCore.Qt.Horizontal and role in (
                QtCore.Qt.DisplayRole, QtCore.Qt.EditRole):
            try:
                self.horizontalHeaders[section] = data
                return True
            except:
                return False
        return super().setHeaderData(section, orientation, data, role)

    def headerData(self, section, orientation, role=QtCore.Qt.DisplayRole):
        if orientation == QtCore.Qt.Horizontal and role == QtCore.Qt.DisplayRole:
            try:
                return self.horizontalHeaders[section]
            except:
                pass
        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=QtCore.Qt.DisplayRole):
        if index.isValid():
            if role == QtCore.Qt.DisplayRole:
                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))
        return QtCore.QVariant()


class MainWindow(QMainWindow):
    def __init__(self, ui_file: str, categories: typing.Sequence):
        super(MainWindow, self).__init__()
        uic.loadUi(ui_file, self)

        self._createCategoryButton \
            = self.findChild(QPushButton, "createCategoryButton")
        self._deleteCategoryButton \
            = self.findChild(QPushButton, "deleteCategoryButton")
        self._applyCategoryButton \
            = self.findChild(QPushButton, "applyCategoryButton")

        self._categoryListView = self.findChild(QListView, "categoryListView")
        self._statementTableView = self.findChild(QTableView,
                                                  "statementTableView")

        self._set_categories(categories)

        self._statementTableView.setSelectionBehavior(QTableView.SelectRows)

    def _set_categories(self, categories: typing.Sequence[str]):
        model = QtGui.QStandardItemModel()
        self._categoryListView.setModel(model)

        for category in categories:
            item = QtGui.QStandardItem(category)
            model.appendRow(item)

    def set_statement_data(self, df: pd.DataFrame):
        model = PandasModel(df)
        self._statementTableView.setModel(model)


def show_main_window(ui_file, categories: typing.Sequence[str] = None,
                     df: pd.DataFrame = None):
    if categories is None:
        categories = []

    app = QApplication(sys.argv)
    window = MainWindow(ui_file, categories)

    if df is not None:
        window.set_statement_data(df)

    window.show()
    app.exec()


def main():
    from banking_breakdown.statement_parser import get_stripped_statement

    categories = ["Food", "Rent & Utilities", "Hobbies", "Education",
                  "Transportation", "Social Life", "Other"]

    df = get_stripped_statement("../res/banking_statement_2023.csv")

    show_main_window("../res/main_window.ui", categories, df)


if __name__ == "__main__":
    main()