import * as XLSX from 'xlsx';
import { saveAs } from 'file-saver';

type DataType = { [key: string]: string | number | null };

type Column = {
  title: string;
  dataIndex: string;
  key: string;
  children?: Column[];
};

type Merge = {
  s: { r: number; c: number };
  e: { r: number; c: number };
};

function getMaxDepth(columns: Column[]): number {
  let maxDepth = 0;
  for (const column of columns) {
    if (column.children) {
      const depth = 1 + getMaxDepth(column.children);
      maxDepth = Math.max(maxDepth, depth);
    } else {
      maxDepth = Math.max(maxDepth, 1);
    }
  }
  return maxDepth;
}

function processColumns(columns: Column[], startRow = 0, startCol = 0, maxDepth = 0, colLength = 0): { headers: any[][], merges: Merge[] } {
  let headers: any[][] = Array.from({ length: maxDepth }, () => Array(colLength).fill(''));
  let merges: Merge[] = [];
  let colIndex = startCol;

  for (const column of columns) {
    let endRow = startRow;
    let endCol = colIndex;

    if (column.children) {
      const result = processColumns(column.children, startRow + 1, colIndex, maxDepth, colLength);

      headers = headers.map((row, i) => {
        if (i === startRow) {
          return [...row.slice(0, colIndex), ...result.headers[i], ...row.slice(colIndex + result.headers[i].length)];
        } else {
          return [...row.slice(0, colIndex), ...result.headers[i].slice(colIndex, result.headers[i].length), ...row.slice(colIndex + result.headers[i].length)];
        }
      });

      merges = [...merges, ...result.merges];
      endCol += getAllColumns(column.children).length - 1;
      endRow = startRow;

    } else {
      endRow = maxDepth - 1;
    }

    for (let i = startRow; i <= endRow; i++) {
      headers[i][colIndex] = i === startRow ? column.title : '';
    }

    merges.push({ s: { r: startRow, c: colIndex }, e: { r: endRow, c: endCol } });
    colIndex = endCol + 1;
  }

  return { headers, merges };
}

function getAllColumns(columns: Column[]): Column[] {
  let allColumns: Column[] = [];
  for (const column of columns) {
    if (column.children) {
      allColumns = [...allColumns, ...getAllColumns(column.children)];
    } else {
      allColumns.push(column);
    }
  }
  return allColumns;
}

function exportTableToMultiExcel(sheetsData:any[], fileName = 'data.xlsx', needSummary = false) {
  const wb = XLSX.utils.book_new();

  sheetsData.forEach(({ data, columns, sheetName }) => {
    const allColumns = getAllColumns(columns);

    // 计算总结信息
    const summary: DataType = {};
    if (needSummary) {
      allColumns.forEach(column => {
        summary[column.dataIndex] = data.reduce((total: number, row: { [x: string]: number; }) => {
          if (typeof row[column.dataIndex] === 'number') {
            return total + (row[column.dataIndex] as number);
          }
          return total;
        }, 0);
      });
      summary[allColumns[0].dataIndex] = '总计';
    }

    const dataValues = [...data.map((row: { [x: string]: any; }) => {
      const rowData = [];
      for (const column of allColumns) {
        rowData.push(row[column.dataIndex]);
      }
      return rowData;
    }), needSummary ? Object.values(summary) : []];  // 如果需要汇总，在数据数组末尾添加总结信息

    const maxDepth = getMaxDepth(columns);
    const { headers, merges } = processColumns(columns, 0, 0, maxDepth, allColumns.length);

    const ws = XLSX.utils.aoa_to_sheet([...headers, ...dataValues]);

    const colWidths = allColumns.map((column, i) => {
      let maxLength = column.title.length * 2;
      dataValues.forEach(row => {
        const cellValue = row[i];
        if (cellValue) {
          const cellLength = cellValue.toString().length;
          const cellWidth = /[\u4e00-\u9fa5]/.test(cellValue.toString()) ? cellLength * 2 : cellLength;
          maxLength = Math.max(maxLength, cellWidth);
        }
      });
      return { wch: maxLength };
    });

    ws['!cols'] = colWidths;

    const range = XLSX.utils.decode_range(ws['!ref'] as string);
    for (let R = range.s.r; R <= range.e.r; ++R) {
      for (let C = range.s.c; C <= range.e.c; ++C) {
        const cell_address = { c: C, r: R };
        const cell_ref = XLSX.utils.encode_cell(cell_address);
        if (!ws[cell_ref]) ws[cell_ref] = {};
        ws[cell_ref].s = {
          alignment: {
            horizontal: 'center',
            vertical: 'center'
          }
        };
      }
    }

    for (let r = 0; r < maxDepth; r++) {
      for (let c = 0; c < headers[r].length; c++) {
        if (!ws[XLSX.utils.encode_cell({ r, c })]) {
          ws[XLSX.utils.encode_cell({ r, c })] = {};
        }
        ws[XLSX.utils.encode_cell({ r, c })].s = {
          fill: {
            fgColor: { rgb: "FFFF00" }
          },
          alignment: {
            horizontal: 'center',
            vertical: 'center'
          }
        };
      }
    }

    ws['!merges'] = merges;
    XLSX.utils.book_append_sheet(wb, ws, sheetName);
  });

  const wbout = XLSX.write(wb, { bookType: 'xlsx', type: 'array' });
  const blob = new Blob([wbout], { type: 'application/octet-stream' });
  saveAs(blob, fileName);
}

export default exportTableToMultiExcel;
