import sys
import pandas as pd
from PyQt5.QtCore import QAbstractTableModel, QModelIndex, Qt
from PyQt5.QtWidgets import QApplication, QTableView, QStyledItemDelegate, QVBoxLayout, QWidget, QPushButton, QLineEdit, QHeaderView
from PyQt5.QtCore import QItemSelection, QItemSelectionModel
from PyQt5.QtGui import QPainter, QColor, QFont
from PyQt5.QtCore import Qt, QModelIndex
from PyQt5.QtCore import QVariant


from .theme_manager import ThemeManager

# Custom Header View
class CustomHeaderView(QHeaderView):
    def __init__(self, orientation, parent=None, mainparent=None):
        super().__init__(orientation, parent)
        self.setDefaultAlignment(Qt.AlignCenter)
        self.mainparent = mainparent 

    def paintSection(self, painter, rect, logicalIndex):
        if not rect.isValid():
            print("rect not valid")
            return

        # Explicitly fill the background to prevent transparency
        # painter.fillRect(rect, QColor(240, 240, 240))  # Light gray background

        # Retrieve the header text
        model = self.model()
        text = model.headerData(logicalIndex, self.orientation(), Qt.DisplayRole)

        if text and '\n' in text:
            first_line, second_line = text.split('\n', 1)
        else:
            first_line, second_line = text, ''

        painter.save()
        try:
            painter.setRenderHint(QPainter.TextAntialiasing)

            # Paint the background
            theme_manager = ThemeManager.instance
            if theme_manager.theme == "bright":
                painter.fillRect(rect, QColor(240, 240, 240))  # Light gray background
            else:
                painter.fillRect(rect, QColor(34, 34, 34))  # Light gray background

            # Set text font and colors
            font = painter.font()
            font.setBold(True)
            painter.setFont(font)

            # Paint the first line in a custom color
            #painter.setPen(QColor(0, 128, 255))  # Blue color
            try:
                painter.setPen(self.mainparent.theme_manager.get_color(7))
            except:
                try:
                    painter.setPen(self.mainparent.parent().theme_manager.get_color(7))
                except:
                    painter.setPen(QColor(150,150,150))

            painter.drawText(rect, Qt.AlignTop | Qt.AlignLeft, first_line)

            # Paint the second line in black
            if theme_manager.theme == "bright": 
                painter.setPen(Qt.black)
            else:
                painter.setPen(Qt.white)
            painter.drawText(rect, Qt.AlignBottom | Qt.AlignLeft, second_line)
        finally:
            painter.restore()


class PandasModel(QAbstractTableModel):

    def __init__(self, df, view=None, sorting=None, blocked_cols=None):
        super(PandasModel, self).__init__()

        if sorting:
            df = df[sorting]

        if blocked_cols is None:
            self.blocked_cols = []
        else:
            self.blocked_cols = blocked_cols

        self._df = df
        self.view = view
        self.set_data(df)

        self.df0 = df.copy()

        self.last_search_index = None 

        self.first_init = True 

    #def setModel(self, model):
    #    super().setModel(model)
    @property 
    def df(self):
        return self._df
    
    def rowCount(self, parent=QModelIndex()):
        return len(self._df)

    def columnCount(self, parent=QModelIndex()):
        return len(self._df.columns)

    def _is_number(self, value):
        """Return True if value is numeric or numeric-like string."""
        try:
            float(str(value).replace(',', ''))
            return True
        except ValueError:
            return False

    def data(self, index, role=Qt.DisplayRole):
        if self.view is None:
            is_selected = False
            is_editing = False 
        else:
            is_selected = self.view.selectionModel().isSelected(index)
            is_editing = self.view.state() == QTableView.EditingState and self.view.currentIndex() == index


        if role == Qt.DisplayRole and not is_editing:
            value = self._df.iloc[index.row(), index.column()]
            return str(value)
        elif role == Qt.DisplayRole or is_editing:
            return None 
        elif role == Qt.EditRole:
            value = self._df.iloc[index.row(), index.column()]
            return str(value)
        
        if role == Qt.TextAlignmentRole:
            value = self._df.iloc[index.row(), index.column()]
            if self._is_number(value):
                return Qt.AlignRight | Qt.AlignVCenter
            else:
                return Qt.AlignLeft | Qt.AlignVCenter
        return None
    

    def setData(self, index, value, role=Qt.EditRole):

        if self.view is not None:
            header = self.view.horizontalHeader()
            header.setDefaultAlignment(Qt.AlignLeft | Qt.AlignVCenter)

        if role == Qt.EditRole:
            self._df.iloc[index.row(), index.column()] = value
            self.dataChanged.emit(index, index, [Qt.DisplayRole])
            return True
        return False
    
    def flags(self, index):
        if len(self.blocked_cols) > 0:
            for block in self.blocked_cols:
                column_index = self._df.columns.get_loc(block)
                if index.column() == column_index:
                    return Qt.ItemIsEnabled | Qt.ItemIsSelectable
        return Qt.ItemIsEditable | Qt.ItemIsEnabled | Qt.ItemIsSelectable

    # def headerData(self, section, orientation, role=Qt.DisplayRole):
    #     if orientation == Qt.Horizontal and role == Qt.DisplayRole:
    #         return str(self._df.columns[section])
    #     if orientation == Qt.Vertical and role == Qt.DisplayRole:
    #         return str(self._df.index[section])
    #     return None

    def headerData(self, section, orientation, role=Qt.DisplayRole):
        if role == Qt.DisplayRole:
            if orientation == Qt.Horizontal:
                return str(self._df.columns[section])
            elif orientation == Qt.Vertical:
                return str(self._df.index[section])
        return QVariant()


    def addRow(self, row_data=None):
        """Adds a new row to the model.
        
        Parameters:
        row_data: list of data to insert in the new row. Default is None (will create an empty row with NA values).        
        """
        # Get the selection model to evaluate selected cells or rows
        try:
            selection_model = self.view.selectionModel()
            
            if selection_model.hasSelection():
                # Get selected indexes (QModelIndex objects)
                selected_indexes = selection_model.selectedIndexes()
                
                if selected_indexes:
                    # Find the highest row index from the selected cells
                    max_row_index = max(index.row() for index in selected_indexes)
                    row_index = max_row_index + 1  # Insert below the highest selected row
                else:
                    # Insert at the end if no valid selection is made
                    row_index = self.rowCount()
            else:
                # Insert at the end if no selection is made
                row_index = self.rowCount()
        except Exception as e:
            print("Exception:", str(e))
            row_index = self.rowCount()

        self.beginInsertRows(QModelIndex(), row_index, row_index)

        if row_data is None:
            print("no row data")
            row_data = [""] * len(self._df.columns)  # Create a list of NA values matching the number of columns
            new_row_df = pd.DataFrame([row_data], columns=self._df.columns)
            self._df = pd.concat([self._df.iloc[:row_index], new_row_df, self._df.iloc[row_index:]]).reset_index(drop=True)
        
        else:
            # print("add row data", row_data)
            if len(row_data) != len(self._df.columns):
                print("Exception: row data was length %s but require %s" % (len(row_data), len(self._df.columns)))
                return
            new_row_df = pd.DataFrame([row_data], columns=self._df.columns)
            self._df = pd.concat([self._df.iloc[:row_index], new_row_df, self._df.iloc[row_index:]]).reset_index(drop=True)
        
        self._df = self._df.replace(pd.NA, "")
        self.endInsertRows()


    def removeSelectedRows(self):
        """Remove selected rows."""
        selected_rows = self.view.selectionModel().selectedRows()  # Get selected row indices
        rows_to_remove = sorted([index.row() for index in selected_rows], reverse=True)  # Sort in reverse to avoid index shifting

        for row_index in rows_to_remove:
            self.beginRemoveRows(QModelIndex(), row_index, row_index)
            self._df.drop(self._df.index[row_index], inplace=True)
            self._df.reset_index(drop=True, inplace=True)  # Reset index after dropping rows
            self.endRemoveRows()

        # Emit signal to update view after row removal
        self.layoutChanged.emit()

    def set_data(self, new_data):
        """Replaces the current DataFrame with a new DataFrame."""
        self.beginResetModel()
        self._df = new_data
        self.endResetModel()

        if self.view is not None:
            header = self.view.horizontalHeader()
            header.setDefaultAlignment(Qt.AlignLeft | Qt.AlignVCenter)

    @property 
    def raw_data(self):
        return self._df
    

    def swap_up(self):
        """Swap the currently selected row with the one above, update the selection, and scroll to the new selection."""
        selected_rows = self.view.selectionModel().selectedRows()
        if selected_rows:
            row_index = selected_rows[0].row()  # Assume single selection
            if row_index > 0:  # Can only swap if not at the top
                # Swap the rows in the DataFrame
                self._df.iloc[[row_index, row_index - 1]] = self._df.iloc[[row_index - 1, row_index]].values
                self.layoutChanged.emit()  # Notify the view that the data has changed

                # Update the selection to the new row
                new_selection = self.index(row_index - 1, 0)
                self.view.selectionModel().clearSelection()  # Clear the previous selection
                self.view.selectionModel().select(new_selection, QItemSelectionModel.Select | QItemSelectionModel.Rows)  # Select the new row
                
                # Scroll to the new selection
                self.view.scrollTo(new_selection, QTableView.PositionAtCenter)

    def swap_down(self):
        """Swap the currently selected row with the one below, update the selection, and scroll to the new selection."""
        selected_rows = self.view.selectionModel().selectedRows()
        if selected_rows:
            row_index = selected_rows[0].row()  # Assume single selection
            if row_index < self.rowCount() - 1:  # Can only swap if not at the bottom
                # Swap the rows in the DataFrame
                self._df.iloc[[row_index, row_index + 1]] = self._df.iloc[[row_index + 1, row_index]].values
                self.layoutChanged.emit()  # Notify the view that the data has changed

                # Update the selection to the new row
                new_selection = self.index(row_index + 1, 0)
                self.view.selectionModel().clearSelection()  # Clear the previous selection
                self.view.selectionModel().select(new_selection, QItemSelectionModel.Select | QItemSelectionModel.Rows)  # Select the new row
                
                # Scroll to the new selection
                self.view.scrollTo(new_selection, QTableView.PositionAtCenter)
    
    def try_focus(self, searchstr: str, columns=None):
        """Search for the next occurrence of searchstr in the DataFrame and focus the selection on it."""
        
        if columns is None:
            columns = list(range(self.columnCount()))

        start_row = 0
        start_col = 0

        # If this is not the first search, start from the next cell
        if self.last_search_index:
            start_row, start_col = self.last_search_index.row(), self.last_search_index.column() + 1
            if start_col >= self.columnCount():
                start_row += 1
                start_col = 0
        
        # Search from the last found position
        for row in range(start_row, self.rowCount()):
            for col in range(start_col, self.columnCount()):
                cell_value = str(self._df.iloc[row, col])
                if searchstr in cell_value:  # Case-sensitive search
                    # Match found, update selection
                    index = self.index(row, col)
                    self.view.selectionModel().clearSelection()  # Clear previous selection
                    self.view.selectionModel().select(index, QItemSelectionModel.Select | QItemSelectionModel.Rows)  # Select the found row

                    # Scroll to the found row
                    self.view.scrollTo(index, QTableView.PositionAtCenter)

                    # Store the last found index for the next search
                    self.last_search_index = index
                    return  # Stop searching after the first match

            # Reset the start column for subsequent rows
            start_col = 0

        # If no match found, reset the search index for future searches
        print(f"'{searchstr}' not found or no more occurrences.")
        self.last_search_index = None
    
    def _apply_column_order(self, df: pd.DataFrame) -> pd.DataFrame:
        """Return df with columns reordered per self.order_cols.
        Any missing/extra columns are handled gracefully."""
        if not self.order_cols:
            return df

        # Keep only columns that exist; preserve requested order
        requested = [c for c in self.order_cols if c in df.columns]

        # Append any columns not mentioned, preserving their current order
        remaining = [c for c in df.columns if c not in requested]
        new_order = requested + remaining

        # No-op if identical
        if new_order == list(df.columns):
            return df

        return df[new_order]

    def _apply_row_order(self, df: pd.DataFrame) -> pd.DataFrame:
        """Return df with rows reordered/sorted per self.order_row.

        Supports:
          - list/tuple: explicit order (by index labels or integer positions)
          - str: sort_values(by=str, ascending=True)
          - dict: passed to sort_values(**dict)
        """
        if self.order_row is None:
            return df

        spec = self.order_row

        # Explicit order by index labels or by integer positions
        if isinstance(spec, (list, tuple)):
            if len(spec) == 0:
                return df

            first = spec[0]
            # If entries look like integer positions
            if all(isinstance(x, int) for x in spec):
                # clamp to valid positions
                pos = [p for p in spec if 0 <= p < len(df)]
                # reindex by position, then append leftovers in original order
                ordered = df.iloc[pos]
                leftovers_mask = ~df.index.isin(ordered.index)
                leftovers = df.loc[leftovers_mask]
                return pd.concat([ordered, leftovers], axis=0)
            else:
                # treat entries as index *labels*
                existing_labels = [x for x in spec if x in df.index]
                ordered = df.loc[existing_labels]
                leftovers = df.drop(index=existing_labels, errors='ignore')
                return pd.concat([ordered, leftovers], axis=0)

        # Sort by column name (ascending)
        if isinstance(spec, str):
            if spec in df.columns:
                return df.sort_values(by=spec, ascending=True, kind="mergesort")  # stable
            return df  # unknown column → no-op

        # Dict → passthrough to sort_values
        if isinstance(spec, dict):
            # Accept keys like by, ascending, kind, key, na_position, etc.
            kwargs = dict(spec)
            by = kwargs.get("by")
            if by is None:
                return df
            # validate columns
            valid_by = [c for c in (by if isinstance(by, (list, tuple)) else [by]) if c in df.columns]
            if not valid_by:
                return df
            kwargs["by"] = valid_by
            # ensure stable default if not specified
            kwargs.setdefault("kind", "mergesort")
            try:
                return df.sort_values(**kwargs)
            except Exception:
                return df  # fail safe

        # Unknown type → no-op
        return df

    def rearrange(self):
        """Reorder columns/rows based on order_cols/order_row and refresh the view."""
        self.beginResetModel()
        df = self._df

        # Columns first (affects column positions used by row operations that reference columns)
        df = self._apply_column_order(df)

        # Then rows
        df = self._apply_row_order(df)

        self._df = df.reset_index(drop=False) if isinstance(self.order_row, (list, tuple)) and not all(isinstance(x, int) for x in self.order_row) and df.index.dtype != 'int64' else df
        # ^ keeps things robust when explicit label reordering leaves non-range index. If you don't
        #   want index reset, you can drop the reset_index line and keep original indices.

        self.endResetModel()

        if self.view is not None:
            header = self.view.horizontalHeader()
            header.setDefaultAlignment(Qt.AlignLeft | Qt.AlignVCenter)

    def set_order(self, cols=None, rows=None, apply_now=True):
        """Convenience setter for ordering."""
        if cols is not None:
            self.order_cols = list(cols)
        if rows is not None:
            self.order_row = rows
        if apply_now:
            self.rearrange()