import React from 'react';
import cx from 'classnames';
import type { TableExtension } from '../../base/types';

const CLASS_NAME = 'TableColumnSpan';

export const TableColumnSpanExtension: TableExtension = {
    name: 'columnSpan',
    isEnabled: (props) => !!props.columnSpan,
    onPropsInit: (props) => ({
        ...props,
        columnDefs: props.columnDefs.map((def) => ({
            ...def,
            rowSpan: (row) => {
                const spanFlag = row.flags[`${CLASS_NAME}-span`];
                const count = typeof spanFlag === 'boolean' ? 0 : spanFlag;

                if (props.columnSpan?.columns.includes(def.name)) return count;
                return 1;
            },
        })),
    }),
    onRowsCreated: (rows, { columnSpan }) => {
        if (!columnSpan || rows.length === 0) return rows;

        // Create groups of adjacent rows which share a spanKey
        let current: { key: string; count: number } | undefined;
        const groups = rows
            .reduce<{ key: string; count: number }[]>((groups, row, index) => {
                const key = columnSpan.spanKey(row.data).toString();
                let returnValue = groups;

                if (current && current.key === key) {
                    current.count++;
                }

                if (current?.key !== key) {
                    if (current) returnValue = [...returnValue, current];
                    current = { key, count: 1 };
                }

                if (current && index === rows.length - 1) {
                    returnValue = [...returnValue, current];
                }

                return returnValue;
            }, [])
            .map((group, index) => {
                return Array(group.count)
                    .fill(group)
                    .map(({ count }: { count: number }, i) => ({ count, primary: i === 0, index }));
            })
            .flat();

        return rows.map((row, index) => {
            const group = groups[index];
            return {
                ...row,
                flags: {
                    [`${CLASS_NAME}-span`]: group.primary ? group.count : 0,
                    [`${CLASS_NAME}-odd`]: group.index % 2 === 0,
                    [`${CLASS_NAME}-even`]: group.index % 2 === 1,
                },
            };
        });
    },

    rowWrapper: (rows, { columnSpan }) => {
        if (!columnSpan) return rows;

        // Wrap span groups in a div to support group styling
        const groups: typeof rows = [];
        const pushGroup = (groupedRows: typeof rows) => {
            if (groupedRows.length === 1) return groups.push(groupedRows[0]);
            return groups.push([
                <div
                    className={cx(
                        `${CLASS_NAME}-group`,
                        `${CLASS_NAME}-group-align-${columnSpan.alignment ?? 'top'}`
                    )}
                    key={columnSpan.spanKey(groupedRows[0][1].data)}
                >
                    {groupedRows.map(([element]) => element)}
                </div>,
                groupedRows[0][1],
            ]);
        };

        let prev: typeof rows = [];
        let prevKey: null | string = null;
        rows.forEach((row) => {
            const key = columnSpan.spanKey(row[1].data).toString();
            if (prevKey === null || key === prevKey) prev.push(row);
            else {
                pushGroup(prev);
                prev = [row];
            }
            prevKey = key;
        });
        if (prev.length > 0) pushGroup(prev);

        return groups;
    },
};