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

const CLASS_NAME = 'TableGroupedRows';

export const TableGroupedRowsExtension: TableExtension = {
    name: 'groupedRows',
    isEnabled: (props) => !!props.grouped,
    onPropsInit: (props) => ({
        ...props,
        rowSelectEnabled: (row) => {
            if (
                props.grouped?.rowSpanColumns?.includes('table-selection') &&
                !row.flags[`${CLASS_NAME}-primary`]
            )
                return false;
            return props?.rowSelectEnabled?.(row) ?? true;
        },
        columnDefs: props.columnDefs.map((def) => ({
            ...def,
            rowSpan: (row) => {
                const spanFlag = row.flags[`${CLASS_NAME}-span`];
                const count = typeof spanFlag === 'boolean' ? 0 : spanFlag;

                if (!row.flags[`${CLASS_NAME}-expanded`]) return count;

                if (props.grouped?.rowSpanColumns?.includes(def.name)) return count;
                return 1;
            },
        })),
        onSorted: (rows) => {
            // Pull rows up to the top row from that group to ensure they're still grouped together
            const sorted = props.onSorted?.(rows) ?? rows;
            const { grouped } = props;
            if (!grouped) return sorted;

            const groupIndex: Record<string, number> = {};
            const sortedGroups: (typeof rows)[] = [];
            sorted.forEach((row) => {
                const groupKey = grouped?.aggregationKey(row);
                const index = groupIndex[groupKey] ?? sortedGroups.length;
                groupIndex[groupKey] = index;
                if (!sortedGroups[index]) sortedGroups[index] = [];
                sortedGroups[index].push(row);
            });
            return sortedGroups.flat();
        },
    }),
    onRowsCreated: (rows, { grouped }) => {
        if (!grouped) return rows;

        const groups = rows.reduce<Record<string, { count: number; primary: number }>>(
            (groups, row) => {
                const key = grouped.aggregationKey(row.data);
                const group = !groups[key]
                    ? { primary: row.index, count: 1 }
                    : { ...groups[key], count: groups[key].count + 1 };
                return {
                    ...groups,
                    [key]: group,
                };
            },
            {}
        );

        let visibleCount = -1;
        return rows.map((row) => {
            const groupKey = grouped.aggregationKey(row.data);
            const group = groups[groupKey];
            const isPrimary = group.primary === row.index;
            const isExpanded = grouped.expandedGroups[groupKey];
            if (isPrimary || isExpanded) visibleCount++;
            return {
                ...row,
                flags: {
                    [`${CLASS_NAME}-span`]: isPrimary ? group.count : 0,
                    [`${CLASS_NAME}-expanded`]: isExpanded,
                    [`${CLASS_NAME}-collapsed`]: !isExpanded,
                    [`${CLASS_NAME}-primary`]: isPrimary,
                    [`${CLASS_NAME}-isGroup`]: group.count > 1,
                    [`${CLASS_NAME}-odd`]: visibleCount % 2 == 0,
                    [`${CLASS_NAME}-even`]: visibleCount % 2 == 1,
                },
            };
        });
    },

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

        const groups: typeof rows = [];
        let prev: typeof rows = [];

        const pushGroup = (groupedRows: typeof rows) => {
            if (groupedRows.length === 1) return groups.push(groupedRows[0]);
            return groups.push([
                <div
                    className={`${CLASS_NAME}-group`}
                    key={grouped.aggregationKey(groupedRows[0][1].data)}
                >
                    {groupedRows.map(([element]) => element)}
                </div>,
                groupedRows[0][1],
            ]);
        };

        rows.forEach((row, index) => {
            const key = grouped.aggregationKey(row[1].data);
            const prevKey = index !== 0 ? grouped.aggregationKey(rows[index - 1][1].data) : null;
            if (prevKey === null || key === prevKey) prev.push(row);
            else {
                pushGroup(prev);
                prev = [row];
            }
        });
        pushGroup(prev);

        return groups;
    },
};