geotech_db_sticks_table = editable_table(
    dataframe=GroundModelManagerState.geotech_db_sticks_df,
    columns=GroundModelManagerState.geotech_db_sticks_df_columns,
    on_selection_change=GroundModelManagerState.submit_selected_rows,
    has_title=True,
    title="Sticks",
    has_searchbar=True,
)

from dataclasses import dataclass
from typing import Callable
from typing import Any
import pandas as pd
import polars as pl

from reflex.components.datadisplay.dataeditor import DataEditorTheme
import reflex as rx

beca_theme = {"accent_color": "teal"}


@dataclass
class CRUDOperation:
    """ Class to represent CRUD operations for tracking changes in the editable table."""

    type: str
    column_index: int = None
    row_index: int = None
    old_row: dict = None
    new_row: dict = None


def editable_table(
    dataframe: pd.DataFrame | pl.DataFrame,
    columns: list[dict[str, str]],
    on_save_edits: Callable | list[Callable] = None,
    on_submit_selection: Callable | list[Callable] = None,
    on_selection_change: Callable | list[Callable] = None,
    title: str = "",
    has_title: bool = False,
    has_searchbar: bool = False,
    **props
):
    """ A reusable editable table component with custom theming.

    Notes:
    - This will scale up to whatever sized box it is placed in, if the box does not have a set size, it will beloon vertically instead of creating a scroll.
    """

    if isinstance(dataframe, pl.DataFrame):
        dataframe = dataframe.to_pandas()

    return EditableTable.create(
        dataframe=dataframe,
        columns=columns,
        on_save_edits=on_save_edits,
        on_submit_selection=on_submit_selection,
        on_selection_change=on_selection_change,
        title=title,
        has_title=has_title,
        has_searchbar=has_searchbar,
        **props
    )


class EditableTable(rx.ComponentState):
    """ A reusable editable table component with custom theming.

    Notes:
    - This will scale up to whatever sized box it is placed in, if the box does not have a set size, it will beloon vertically instead of creating a scroll.
    """
    raw_columns: list[dict[str, str]] = []

    raw_data: list[list[Any]] = []

    changes: list[CRUDOperation] = []

    is_editing: bool = False

    search_value: str = ""

    current_grid_selection: dict = None

    @rx.event
    def update_internal_data(self, external_dataframe_dict: dict[str, Any], columns: list[dict[str, str]]):
        """Refreshes the table from external dataframe (Reflex auto-converts to dict format sing serializer)."""

        if not external_dataframe_dict or not isinstance(external_dataframe_dict, dict):
            print("No external dataframe provided to update internal data.")
            return

        raw_columns = columns

        raw_data = external_dataframe_dict.get("data", [])

        if self.raw_data != raw_data or self.raw_columns != raw_columns:
            self.raw_columns = raw_columns
            self.raw_data = raw_data

    @rx.var
    def cleaned_data(self) -> list[list]:
        if not self.raw_columns or not self.raw_data:
            return []

        # Ensure raw_data is a list and handle None case
        if self.raw_data is None or not isinstance(self.raw_data, list):
            return []

        data = self.raw_data.copy()

        for row in data:
            # Ensure row is a list
            if not isinstance(row, list):
                continue

            # Go through each cell and check for null or nan and convert them to empty string
            for cell_idx, cell in enumerate(row):
                if cell is None:
                    row[cell_idx] = ""
                elif isinstance(cell, (int, float)) and pd.isna(cell):
                    row[cell_idx] = ""

        return data

    @rx.var
    def edited_data(self) -> list[list]:
        """Return the current edited data."""

        # Take the list of CrudOperations and apply them to the cleaned data to get the current edited data.
        data = self.cleaned_data.copy()
        changes = self.changes

        if not changes:
            return data

        for change in changes:
            if change.type == "update":
                data[change.row_index][change.column_index] = change.new_row[
                    list(change.new_row.keys())[change.column_index]
                ]
            elif change.type == "create":
                data.append(change.new_row)
            elif change.type == "delete":
                data.pop(change.row_index)

        return data

    @rx.var
    def filtered_data(self) -> list[list]:
        """Return filtered data based on search value."""

        data = self.edited_data
        search_value = self.search_value

        if not search_value or not data:
            return data

        filtered = []
        for row in data:
            row_matches = False
            for cell_value in row:
                if isinstance(cell_value, bool):
                    continue

                cell_str = (
                    str(cell_value).lower()
                    if cell_value is not None else ""
                )

                if search_value in cell_str:
                    row_matches = True
                    break

            if row_matches:
                filtered.append(row)

        return filtered

    @rx.var
    def display_columns(self) -> list[dict[str, str]]:
        """ Return the columns to be displayed, either the working columns in edit mode or the converted raw columns in view mode."""
        # Ensure we always return a valid list, even if empty
        return self.raw_columns if self.raw_columns else []

    @rx.var
    def display_data(self) -> list[list[Any]]:
        """ Return the data to be displayed, either the filtered internal data in edit mode or the filtered raw data in view mode."""
        # Ensure we always return a valid list, even if empty
        return self.filtered_data if self.filtered_data else []

    @rx.var
    def is_loading(self) -> bool:
        """ Returns True if the table is loading (no data yet)."""
        return not bool(self.raw_data) or not bool(self.raw_columns)

    @rx.event
    def on_finished_editing(self, position, data):
        """ Event handler for when editing a cell is finished in the data editor. Creates and saves CRUD operation without modifying raw_data."""

        if not self.is_editing:
            print("Table is not in editing mode. Ignoring edit.")
            return

        column_index = position[0]
        row_index = position[1]
        column_ids = [col["id"] for col in self.raw_columns]
        new_value = data["data"]

        # Get the old row as dict
        old_row_dict = dict(zip(column_ids, self.raw_data[row_index]))

        # Create new row dict with the change
        new_row_dict = old_row_dict.copy()
        new_row_dict[column_ids[column_index]] = new_value

        # Track the change (raw_data will be computed in edited_data var)
        change_record = CRUDOperation(
            type="update",
            column_index=column_index,
            row_index=row_index,
            old_row=old_row_dict,
            new_row=new_row_dict,
        )

        self.changes.append(change_record)

    @rx.event
    def enter_edit_mode(self):
        """ Enter edit mode"""
        if self.is_editing:
            return

        self.changes = []
        self.is_editing = True

    @rx.event
    def exit_edit_mode(self):
        """ Exit edit mode directly."""
        if not self.is_editing:
            return

        self.changes = []
        self.is_editing = False

    @rx.event
    def update_grid_selection(self, selection: dict):
        """ Update the list of selected rows based on grid selection change event."""
        self.current_grid_selection = selection
        print("Grid selection updated: ", selection)

    @rx.event
    def clear_changes(self):
        """ Clear the list of changes made during editing."""
        self.changes = []
        self.is_editing = False
        print("Cleared changes list.")

    @classmethod
    def toggle_button(cls, on_save_edits: Callable | None) -> rx.Component:
        """Render the toggle edit/save button."""

        return rx.cond(
            cls.is_editing,
            rx.button(
                "Save",
                on_click=[on_save_edits(cls.changes), cls.exit_edit_mode] if on_save_edits else [
                    cls.exit_edit_mode]
            ),
            rx.button(
                "Edit",
                on_click=cls.enter_edit_mode
            )
        )

    @rx.event
    def add_row(cls):
        """Add a new empty row to the table."""
        if not cls.is_editing:
            return

        # Create a new row with empty strings matching the number of columns
        num_columns = len(cls.raw_columns)
        new_row = [""] * num_columns

        # Track as Create change
        change = CRUDOperation(
            type="create",
            new_row=new_row
        )
        cls.changes.append(change)

    @rx.event
    def update_search_value(cls, search_value: str):
        # set the search value
        cls.search_value = str(search_value).strip().lower()

    @rx.var
    def currently_selected_rows(self) -> list:
        data = self.filtered_data
        selection = self.current_grid_selection

        if selection is None or "rows" not in selection:
            return []

        print("\nSelection: ", selection)

        rows: dict[list] = selection["rows"]

        items = rows.get("items", [])
        if not items or not isinstance(items, list) or len(items) == 0:
            return []

        first_item = items[0]
        if not isinstance(first_item, list) or len(first_item) == 0:
            return []

        selected_row_index = first_item[0]
        if not isinstance(selected_row_index, int) or selected_row_index < 0 or selected_row_index >= len(data):
            return []

        selected_rows = []
        for item in items:
            # Handle single row selection [row_index, row_index+1] or range selection [start_row, end_row]
            start_index = item[0]
            end_index = item[1]

            # Add all rows in the range [start_index, end_index]
            for row_index in range(start_index, end_index):
                if isinstance(row_index, int) and 0 <= row_index < len(data):
                    selected_values = dict(
                        zip([col["title"] for col in self.raw_columns], data[row_index]))
                    selected_rows.append(selected_values)

        print("Selected rows:", selected_rows)
        return selected_rows

    @classmethod
    def loading_spinner(cls) -> rx.Component:
        return rx.text("loading")

    @classmethod
    def searchbar(cls, has_searchbar: bool) -> rx.Component:
        return rx.cond(
            has_searchbar,
            rx.input(
                rx.input.slot(rx.icon("search")),
                type="search",
                placeholder="Search...",
                show_search=True,
                radius="medium",
                on_change=lambda value: cls.update_search_value(value),
                on_mount=cls.update_search_value(""),
            ),
            None
        )

    @classmethod
    def title_block(cls, title: str, has_title: bool, has_searchbar: bool) -> rx.Component:
        return rx.cond(
            has_title,
            rx.heading(
                title,
                margin_bottom=rx.cond(has_searchbar, "1rem", "0rem")
            ),
            None
        )

    @classmethod
    def data_grid(cls, on_selection_change, **props) -> rx.Component:
        
        # Build event chain: first update selection, then call external handlers
        selection_events = [cls.update_grid_selection]
        
        if on_selection_change:
            # Add external handlers to the chain (they must not take arguments)
            external_handlers = on_selection_change if isinstance(on_selection_change, list) else [on_selection_change]
            selection_events.extend(external_handlers)

        return rx.cond(
            cls.is_loading,
            cls.loading_spinner(),
            rx.data_editor(
                columns=cls.display_columns,
                data=cls.display_data,
                get_cells_for_selection=True,
                on_cell_edited=cls.on_finished_editing,
                on_grid_selection_change=selection_events,
                grid_selection=cls.current_grid_selection,
                theme=DataEditorTheme(**beca_theme),
                fill_handle=True,
                on_paste=True,
                **props
            ),
        )

    @classmethod
    def button_bar(cls, on_save_edits) -> rx.Component:
        cancel_button = rx.cond(
            cls.is_editing,
            rx.button(
                "Cancel",
                on_click=cls.clear_changes,
                color_scheme="gray",
            ),
            None
        )

        add_row_button = rx.cond(
            cls.is_editing,
            rx.button(
                "Add Row",
                on_click=cls.add_row,
                color_scheme="blue",
            ),
            None
        )

        return rx.flex(
            add_row_button,
            cancel_button,
            cls.toggle_button(on_save_edits),
            gap="0.5rem",
        ),

    @classmethod
    def submit_button(cls, submit_event) -> rx.Component:
        """Render the submit button at the bottom right."""
        if not submit_event:
            return None

        on_click_events = []

        # Pass currently_selected_rows to the submit event
        on_click_events.extend(submit_event if isinstance(
            submit_event, list) else [submit_event(cls.currently_selected_rows)])

        return rx.flex(
            rx.button(
                "Submit",
                on_click=on_click_events,
            ),
            justify_content="end",
            margin_top="1rem",
        )

    @classmethod
    def get_component(
        cls,
        dataframe: pd.DataFrame,
        columns: list[dict[str, str]],
        on_save_edits: Callable | list[Callable] = None,
        on_submit_selection: Callable | list[Callable] = None,
        on_selection_change: Callable | list[Callable] = None,
        title: str = "",
        has_title: bool = False,
        has_searchbar: bool = False,
        **props
    ) -> rx.Component:
        return rx.vstack(
            rx.cond(
                has_title or has_searchbar,
                rx.flex(
                    rx.box(
                        cls.title_block(title, has_title, has_searchbar),
                        cls.searchbar(has_searchbar),
                    ),
                    cls.button_bar(on_save_edits),
                    min_width="30rem",  # Ensures spacing for search bar + button bar
                    width="100%",
                    align_items="flex-end",
                    justify_content=rx.cond(
                        has_title or has_searchbar, "space-between", "end")
                ),
                None
            ),
            cls.data_grid(on_selection_change=on_selection_change, **props),
            cls.submit_button(on_submit_selection),
            width="100%",
            on_mount=cls.update_internal_data(dataframe, columns)
        )