import {
    isValidElement,
    Component,
    createRef,
    cloneElement,
    Children,
    Fragment,
    ReactNode,
    ReactElement,
    ComponentType,
    FC,
    ComponentProps,
} from 'react';
import {
    Announcement,
    Table as AnvilTable,
    TableProps as AnvilTableProps,
    TableRowProps,
    TableColumn,
    TableColumnProps as AnvilTableColumnProps,
    TableCellProps as AnvilTableCellProps,
    TableHeaderCellProps,
    TableHeaderSelectionChangeEvent,
    TableSelectionChangeEvent,
    TablePagerSettings,
} from '@servicetitan/design-system';
import { operators as kendoOperators } from '@progress/kendo-react-grid/dist/npm/filterCommon';
import { GridPDFExport as TablePDFExport } from '@progress/kendo-react-pdf';
import { ExcelExport } from '@progress/kendo-react-excel-export';
import { observer } from 'mobx-react';
import { GRID_COL_INDEX_ATTRIBUTE } from '@progress/kendo-react-grid';

import { AsyncDataSource, IdType } from '@servicetitan/data-query';
import { TableState, Selectable } from './table-state';
import {
    SelectionControlType,
    SelectColumnCell,
    SelectHeaderCell,
} from './select-cell/select-cell';
import memoizeOne from 'memoize-one';

import classNames from 'classnames';
import * as Styles from './table.module.css';
import { action, makeObservable, observable } from 'mobx';
import { useTdProps } from './utils/use-td-props';

const { text: textOperators, ...operators } = kendoOperators;
operators.text = [
    ...textOperators.filter(o => o.operator !== 'isnull' && o.operator !== 'isnotnull'),
];

function isColumnElement<
    T = never,
    TId extends IdType = any,
    P = never,
    PId extends IdType = never,
>(element: ReactNode): element is ReactElement<TableColumnProps<T, TId, P, PId>> {
    if (!isValidElement(element)) {
        return false;
    }

    const elementType = (element as ReactElement).type;

    return elementType === TableColumn || (elementType as any).originalType === TableColumn;
}

export interface TableCellProps<
    T = never,
    TId extends IdType = any,
    P = never,
    PId extends IdType = never,
> extends AnvilTableCellProps {
    tableState?: TableState<T, TId, P, PId>;
    tdProps?: ComponentProps<'td'> & {
        [GRID_COL_INDEX_ATTRIBUTE]?: number;
    };
}

interface TableColumnProps<
    T = never,
    TId extends IdType = any,
    P = never,
    PId extends IdType = never,
> extends AnvilTableColumnProps {
    cell?: ComponentType<TableCellProps<T, TId, P, PId>>;
    children?: ReactElement<TableColumnProps<T, TId, P, PId>>[];
}

type ExcludedTableProps =
    | 'data'
    | 'dataItemKey'
    | 'editField'
    | 'filter'
    | 'onFilterChange'
    | 'group'
    | 'onGroupChange'
    | 'onExpandChange'
    | 'expandField'
    | 'sort'
    | 'onSortChange'
    | 'onHeaderSelectionChange'
    | 'onSelectionChange'
    | 'selectedField'
    | 'pageable'
    | 'pageSize'
    | 'skip'
    | 'onPageChange'
    | 'total'
    | 'selectable'
    | 'pager';

export interface TableProps<
    T,
    TId extends IdType = any,
    P = never,
    PId extends IdType = never,
> extends Omit<AnvilTableProps, ExcludedTableProps> {
    selectable?: boolean;
    exportable?: boolean;
    hideSelectAll?: boolean;
    exportFileName?: string;
    tableState: TableState<T, TId, P, PId>;
    selectionControl?: SelectionControlType;
    pager?: TablePagerSettings;
}

@observer
export class Table<
    T,
    TId extends IdType = any,
    P = never,
    PId extends IdType = never,
> extends Component<TableProps<T, TId, P, PId>> {
    @observable observableScrollable: TableProps<T, TId, P, PId>['scrollable'];
    @observable observableSelectionControl: TableProps<T, TId, P, PId>['selectionControl'];
    @observable observableHideSelectAll: TableProps<T, TId, P, PId>['hideSelectAll'];
    @observable observableSelectable: TableProps<T, TId, P, PId>['selectable'];

    private tableState = this.props.tableState;

    private ref = createRef<HTMLElement>();

    private lastSelection?: Selectable<T>;

    private selectColumnCell: FC<AnvilTableCellProps> = observer((props: AnvilTableCellProps) => (
        <SelectColumnCell
            {...props}
            isRowSelectable={this.tableState.isRowSelectable(props.dataItem)}
            control={this.observableSelectionControl}
        />
    ));

    private customCellMap = new Map<string, FC<TableCellProps<T, TId, P, PId>>>();

    private selectHeaderCell: FC<TableHeaderCellProps> = observer((props: TableHeaderCellProps) => {
        if (this.observableHideSelectAll || this.tableState.selectionLimit !== Infinity) {
            return null;
        }

        return (
            <SelectHeaderCell
                {...props}
                isSomeRowsSelected={
                    this.observableScrollable !== 'virtual'
                        ? this.tableState.isSomePageRowsSelected
                        : this.tableState.isSomeRowsSelected
                }
                control={this.observableSelectionControl}
            />
        );
    });

    // TODO: rid of "memoizeOne" after migration on React.FC
    private withCellTableState = memoizeOne((children: ReactNode) =>
        this.applyToColumns(children, column => {
            const { cell: Cell, field } = column.props;

            return cloneElement(column, {
                cell: this.getOrCreateCellComponentWithTableState(Cell, field),
            });
        })
    );

    constructor(props: TableProps<T, TId, P, PId>) {
        super(props);

        makeObservable(this);

        this.observableScrollable = this.props.scrollable;
        this.observableSelectionControl = this.props.selectionControl;
        this.observableHideSelectAll = this.props.hideSelectAll;
    }

    @action
    componentDidUpdate() {
        this.observableScrollable = this.props.scrollable;
        this.observableSelectionControl = this.props.selectionControl;
        this.observableHideSelectAll = this.props.hideSelectAll;
    }

    rowRender = (row: ReactElement, rowProps: TableRowProps) => {
        const overrides: typeof row.props = {};

        if (
            this.props.selectable &&
            rowProps.rowType === 'data' &&
            !this.tableState.isRowSelectable(rowProps.dataItem)
        ) {
            overrides.className = classNames(row.props.className, Styles.disabled);
        }

        const result = Object.keys(overrides).length ? cloneElement(row, overrides) : row;

        return this.props.rowRender ? this.props.rowRender(result, rowProps) : result;
    };

    applyToColumns(
        children: ReactNode,
        transformer: (
            column: ReactElement<TableColumnProps<T, TId, P, PId>>
        ) => ReactElement<TableColumnProps<T, TId, P, PId>>
    ): ReactNode {
        return Children.map(children, child => {
            if (!isColumnElement<T, TId, P, PId>(child)) {
                return child;
            }

            return child.props.children
                ? cloneElement(transformer(child), {
                      children: this.applyToColumns(
                          child.props.children,
                          transformer
                      ) as typeof child.props.children,
                  })
                : transformer(child);
        });
    }

    handleHeaderSelectionChange = (ev: TableHeaderSelectionChangeEvent) => {
        const checked = ev.syntheticEvent.currentTarget.checked;

        if (this.observableScrollable !== 'virtual') {
            if (checked) {
                this.tableState.selectPage();
            } else {
                this.tableState.deselectPage();
            }
        } else {
            if (checked) {
                this.tableState.selectAll();
            } else {
                this.tableState.deselectAll();
            }
        }
    };

    handleSelectionChange = (ev: TableSelectionChangeEvent) => {
        if (!ev.dataItem) {
            return;
        }
        if (ev.nativeEvent.shiftKey && this.lastSelection) {
            const rows = this.tableState
                .getRowsBetween(this.lastSelection, ev.dataItem)
                .filter(row => this.tableState.isRowSelectable(row));
            const isAsc =
                this.tableState.dataSource?.idSelector?.(this.lastSelection) ===
                this.tableState.dataSource?.idSelector?.(rows[0]);
            const selection = !ev.dataItem.selected;

            if (selection) {
                const toBeSelected: Selectable<T>[] = [];
                for (const row of isAsc ? rows : rows.reverse()) {
                    if (
                        this.tableState.selectedCount + toBeSelected.length >=
                        this.tableState.selectionLimit
                    ) {
                        break;
                    }

                    if (row.selected) {
                        continue;
                    }

                    toBeSelected.push(row);
                }

                this.tableState.setRowsSelection(toBeSelected, selection);
            } else {
                this.tableState.setRowsSelection(rows, selection);
            }
        } else {
            this.tableState.toggleRowSelection(ev.dataItem);
        }

        this.lastSelection = ev.dataItem;
    };

    scrollToBottom = () => {
        if (this.ref.current) {
            const container = this.ref.current.getElementsByClassName(
                'k-grid-content'
            )[0] as HTMLDivElement;
            container.scrollTop = container.scrollHeight - container.clientHeight * 2;
            container.scroll({ top: container.scrollHeight, behavior: 'smooth' });
        }
    };

    render() {
        if (
            this.tableState.dataSource instanceof AsyncDataSource &&
            this.props.scrollable === 'virtual' &&
            this.props.selectable
        ) {
            return (
                <Announcement
                    status="critical"
                    title="AsyncDataSource doesn't support virtual scroll with selection."
                />
            );
        }

        const table = this.table();
        return (
            <Fragment>
                {table}
                {this.props.exportable && (
                    <Fragment>
                        <TablePDFExport
                            ref={this.tableState.setTablePdfExportRef}
                            fileName={sanitizeExportFileName(this.props.exportFileName)}
                        >
                            {this.props.children}
                            {table}
                        </TablePDFExport>
                        <ExcelExport
                            ref={this.tableState.setTableExcelExportRef}
                            fileName={sanitizeExportFileName(this.props.exportFileName)}
                        >
                            {this.props.children}
                        </ExcelExport>
                    </Fragment>
                )}
            </Fragment>
        );
    }

    private getPageable() {
        if (this.props.scrollable === 'virtual' || !this.tableState.pageSize) {
            return false;
        }

        return this.props.pager ?? true;
    }

    private table = () => {
        const { className, filterOperators, selectable, pager, ...props } = this.props;

        return (
            <AnvilTable
                ref={this.ref}
                {...props}
                dataItemKey={this.tableState.rowIdKey}
                data={[...this.tableState.data]}
                rowRender={this.rowRender}
                editField="inEdit"
                filter={this.tableState.filter}
                filterOperators={filterOperators ?? operators}
                onFilterChange={this.tableState.handleFilterChange}
                group={[...this.tableState.group] as AnvilTableProps['group']}
                onGroupChange={this.tableState.handleGroupChange}
                onExpandChange={this.tableState.handleExpandChange}
                expandField="expanded"
                sort={[...this.tableState.sort]}
                onSortChange={this.tableState.handleSortChange}
                onHeaderSelectionChange={this.handleHeaderSelectionChange}
                onSelectionChange={this.handleSelectionChange}
                selectedField="selected"
                pageable={this.getPageable()}
                pageSize={this.tableState.pageSize}
                skip={this.tableState.skip}
                onPageChange={this.tableState.handlePageChange}
                total={this.tableState.filteredCount}
                className={classNames(Styles.table, className)}
                selectable={{ enabled: selectable, mode: 'multiple' }}
            >
                {selectable && (
                    <TableColumn
                        locked
                        field="selected"
                        id="selected"
                        width="50px"
                        filterable={false}
                        resizable={false}
                        reorderable={false}
                        headerSelectionValue={
                            props.scrollable !== 'virtual'
                                ? this.tableState.isAllPageRowsSelected
                                : this.tableState.isAllRowsSelected
                        }
                        headerCell={this.selectHeaderCell}
                        cell={this.selectColumnCell}
                        headerClassName="lh-1-i ta-center-i"
                        className={classNames(Styles.unselectable, 'lh-1-i', 'ta-center-i')}
                    />
                )}
                {this.withCellTableState(props.children)}
            </AnvilTable>
        );
    };

    private getOrCreateCellComponentWithTableState = (
        Cell: any,
        field?: string
    ): FC<TableCellProps<T, TId, P, PId>> | undefined => {
        if (!Cell) {
            return undefined;
        }

        const CellWithTableState: FC<TableCellProps<T, TId, P, PId>> = (
            props: TableCellProps<T, TId, P, PId>
        ) => {
            const tdProps = useTdProps(props);
            return <Cell {...props} tableState={this.tableState} tdProps={tdProps} />;
        };
        CellWithTableState.displayName = 'CellWithTableState';

        if (field) {
            if (this.customCellMap.has(field)) {
                return this.customCellMap.get(field)!;
            }
            this.customCellMap.set(field, CellWithTableState);
        }
        return CellWithTableState;
    };
}

const sanitizeExportFileName = memoizeOne((fileName: string | undefined) =>
    fileName?.replace(/\.+/g, '_')
);
