001/*
002 *  Copyright (c) 2023-2026, Agents-Flex (fuhai999@gmail.com).
003 *  <p>
004 *  Licensed under the Apache License, Version 2.0 (the "License");
005 *  you may not use this file except in compliance with the License.
006 *  You may obtain a copy of the License at
007 *  <p>
008 *  http://www.apache.org/licenses/LICENSE-2.0
009 *  <p>
010 *  Unless required by applicable law or agreed to in writing, software
011 *  distributed under the License is distributed on an "AS IS" BASIS,
012 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013 *  See the License for the specific language governing permissions and
014 *  limitations under the License.
015 */
016package com.agentsflex.core.file2text.extractor.impl;
017
018import com.agentsflex.core.file2text.extractor.FileExtractor;
019import com.agentsflex.core.file2text.source.DocumentSource;
020import org.apache.poi.hssf.usermodel.HSSFWorkbook;
021import org.apache.poi.ss.usermodel.*;
022import org.apache.poi.xssf.usermodel.XSSFWorkbook;
023
024import java.io.IOException;
025import java.io.InputStream;
026import java.util.Collections;
027import java.util.HashSet;
028import java.util.Set;
029
030/**
031 * Excel 文档提取器（.xlsx, .xlsm, .xls, .xltx）
032 * 输出格式：Markdown 表格，支持多 Sheet、公式计算、特殊字符转义
033 */
034public class ExcelExtractor implements FileExtractor {
035
036    private static final Set<String> KNOWN_MIME_TYPES;
037    private static final String MIME_PREFIX = "application/vnd.openxmlformats-officedocument.spreadsheetml";
038    private static final Set<String> SUPPORTED_EXTENSIONS;
039
040    static {
041        Set<String> mimeTypes = new HashSet<>();
042        mimeTypes.add("application/vnd.openxmlformats-officedocument.spreadsheetml.sheet");
043        mimeTypes.add("application/vnd.openxmlformats-officedocument.spreadsheetml.template");
044        mimeTypes.add("application/vnd.ms-excel");
045        mimeTypes.add("application/vnd.ms-excel.sheet.macroEnabled.12");
046        KNOWN_MIME_TYPES = Collections.unmodifiableSet(mimeTypes);
047
048        Set<String> extensions = new HashSet<>();
049        extensions.add("xlsx");
050        extensions.add("xlsm");
051        extensions.add("xls");
052        extensions.add("xltx");
053        SUPPORTED_EXTENSIONS = Collections.unmodifiableSet(extensions);
054    }
055
056    @Override
057    public boolean supports(DocumentSource source) {
058        String mimeType = source.getMimeType();
059        String fileName = source.getFileName();
060
061        if (mimeType != null && KNOWN_MIME_TYPES.contains(mimeType)) {
062            return true;
063        }
064        if (mimeType != null && mimeType.startsWith(MIME_PREFIX)) {
065            return true;
066        }
067        if (fileName != null) {
068            String ext = getExtension(fileName);
069            if (ext != null && SUPPORTED_EXTENSIONS.contains(ext.toLowerCase())) {
070                return true;
071            }
072        }
073        return false;
074    }
075
076    @Override
077    public String extractText(DocumentSource source) throws IOException {
078        StringBuilder text = new StringBuilder();
079
080        try (InputStream is = source.openStream();
081             Workbook workbook = openWorkbook(is, source.getFileName())) {
082
083            FormulaEvaluator evaluator = workbook.getCreationHelper().createFormulaEvaluator();
084            int sheetCount = workbook.getNumberOfSheets();
085
086            for (int i = 0; i < sheetCount; i++) {
087                Sheet sheet = workbook.getSheetAt(i);
088                String sheetName = sheet.getSheetName();
089
090                // Sheet 标题（Markdown 三级标题）
091                text.append("\n### ").append(escapeMarkdown(sheetName)).append("\n\n");
092
093                // 收集所有非空行数据，用于计算最大列数
094                java.util.List<java.util.List<String>> rowsData = new java.util.ArrayList<>();
095                int maxColumns = 0;
096
097                for (Row row : sheet) {
098                    if (isRowEmpty(row)) {
099                        continue;
100                    }
101                    java.util.List<String> rowData = new java.util.ArrayList<>();
102                    short firstCellNum = row.getFirstCellNum();
103                    short lastCellNum = row.getLastCellNum();
104
105                    for (short j = firstCellNum; j < lastCellNum; j++) {
106                        Cell cell = row.getCell(j);
107                        String cellValue = getCellValue(cell, evaluator);
108                        rowData.add(cellValue != null ? cellValue : "");
109                    }
110                    if (!rowData.isEmpty()) {
111                        rowsData.add(rowData);
112                        maxColumns = Math.max(maxColumns, rowData.size());
113                    }
114                }
115
116                if (rowsData.isEmpty()) {
117                    text.append("*Empty Sheet*\n");
118                    continue;
119                }
120
121                // 输出 Markdown 表格
122                outputMarkdownTable(text, rowsData, maxColumns);
123            }
124
125        } catch (Exception e) {
126            throw new IOException("Failed to extract Excel: " + e.getMessage(), e);
127        }
128
129        return text.toString().trim();
130    }
131
132    /**
133     * 输出 Markdown 格式表格
134     */
135    private void outputMarkdownTable(StringBuilder output,
136                                     java.util.List<java.util.List<String>> rowsData,
137                                     int maxColumns) {
138        // 输出表头行（第一行作为表头）
139        java.util.List<String> headerRow = rowsData.get(0);
140        output.append("| ");
141        for (int i = 0; i < maxColumns; i++) {
142            String cell = i < headerRow.size() ? escapeMarkdownCell(headerRow.get(i)) : "";
143            output.append(cell).append(" | ");
144        }
145        output.append("\n| ");
146
147        // 输出分隔行（默认左对齐 :---）
148        for (int i = 0; i < maxColumns; i++) {
149            output.append(":--- | ");
150        }
151        output.append("\n");
152
153        // 输出数据行
154        for (int r = 1; r < rowsData.size(); r++) {
155            java.util.List<String> rowData = rowsData.get(r);
156            output.append("| ");
157            for (int i = 0; i < maxColumns; i++) {
158                String cell = i < rowData.size() ? escapeMarkdownCell(rowData.get(i)) : "";
159                output.append(cell).append(" | ");
160            }
161            output.append("\n");
162        }
163        output.append("\n");
164    }
165
166    /**
167     * 转义单元格内容中的 Markdown 特殊字符
168     * 主要处理：| \ ` * _ [ ] < > 等，避免破坏表格结构
169     */
170    private String escapeMarkdownCell(String content) {
171        if (content == null || content.isEmpty()) {
172            return "";
173        }
174        // 先转义反斜杠，避免重复转义
175        String escaped = content.replace("\\", "\\\\");
176        // 转义管道符（表格分隔符）
177        escaped = escaped.replace("|", "\\|");
178        // 转义其他可能的 Markdown 语法字符（可选，根据需求开启）
179        // escaped = escaped.replace("*", "\\*").replace("_", "\\_").replace("`", "\\`");
180        // 移除换行符，避免破坏表格行结构
181        escaped = escaped.replace("\n", " ").replace("\r", " ");
182        return escaped.trim();
183    }
184
185    /**
186     * 转义 Sheet 名称中的 Markdown 特殊字符（用于标题）
187     */
188    private String escapeMarkdown(String text) {
189        if (text == null) return "";
190        return text.replace("#", "\\#").replace("*", "\\*");
191    }
192
193    private Workbook openWorkbook(InputStream inputStream, String fileName) throws IOException {
194        String ext = fileName != null ? getExtension(fileName) : null;
195
196        if ("xls".equalsIgnoreCase(ext)) {
197            return new HSSFWorkbook(inputStream);
198        } else if ("xlsx".equalsIgnoreCase(ext) || "xlsm".equalsIgnoreCase(ext) || "xltx".equalsIgnoreCase(ext)) {
199            return new XSSFWorkbook(inputStream);
200        }
201
202        try {
203            return new XSSFWorkbook(inputStream);
204        } catch (Exception e) {
205            throw new IOException("Unable to determine Excel format. Please ensure file extension is correct.", e);
206        }
207    }
208
209    private String getCellValue(Cell cell, FormulaEvaluator evaluator) {
210        if (cell == null) {
211            return "";
212        }
213        switch (cell.getCellType()) {
214            case STRING:
215                return cell.getStringCellValue();
216            case NUMERIC:
217                if (DateUtil.isCellDateFormatted(cell)) {
218                    return cell.getDateCellValue().toString();
219                }
220                double numericValue = cell.getNumericCellValue();
221                if (numericValue == Math.floor(numericValue) && !Double.isInfinite(numericValue)) {
222                    return String.valueOf((long) numericValue);
223                }
224                return String.valueOf(numericValue);
225            case BOOLEAN:
226                return String.valueOf(cell.getBooleanCellValue());
227            case FORMULA:
228                try {
229                    CellValue evaluated = evaluator.evaluate(cell);
230                    return formatEvaluatedValue(evaluated);
231                } catch (Exception e) {
232                    return cell.getCellFormula();
233                }
234            case BLANK:
235            case _NONE:
236            default:
237                return "";
238        }
239    }
240
241    private String formatEvaluatedValue(CellValue evaluated) {
242        if (evaluated == null) {
243            return "";
244        }
245        switch (evaluated.getCellType()) {
246            case STRING:
247                return evaluated.getStringValue();
248            case NUMERIC:
249                double num = evaluated.getNumberValue();
250                if (num == Math.floor(num) && !Double.isInfinite(num)) {
251                    return String.valueOf((long) num);
252                }
253                return String.valueOf(num);
254            case BOOLEAN:
255                return String.valueOf(evaluated.getBooleanValue());
256            case ERROR:
257                return "#ERROR:" + evaluated.getErrorValue();
258            default:
259                return "";
260        }
261    }
262
263    private boolean isRowEmpty(Row row) {
264        if (row == null) return true;
265        for (Cell cell : row) {
266            if (cell != null && cell.getCellType() != CellType.BLANK) {
267                // 空行判断使用简化逻辑，不传 evaluator
268                String value;
269                if (cell.getCellType() == CellType.STRING) {
270                    value = cell.getStringCellValue();
271                } else {
272                    value = String.valueOf(cell);
273                }
274                if (value != null && !value.trim().isEmpty()) {
275                    return false;
276                }
277            }
278        }
279        return true;
280    }
281
282    @Override
283    public int getOrder() {
284        return 10;
285    }
286
287    private String getExtension(String fileName) {
288        if (fileName == null || !fileName.contains(".")) return null;
289        int lastDot = fileName.lastIndexOf('.');
290        return fileName.substring(lastDot + 1);
291    }
292}