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 com.agentsflex.core.file2text.util.EncodingDetectUtil;
021import org.jsoup.Jsoup;
022import org.jsoup.nodes.Document;
023import org.jsoup.nodes.Element;
024import org.jsoup.nodes.Node;
025import org.jsoup.select.Elements;
026
027import java.io.IOException;
028import java.io.InputStream;
029import java.util.*;
030import java.util.concurrent.ConcurrentHashMap;
031import java.util.regex.Pattern;
032
033/**
034 * 增强版 HTML 文档提取器
035 * 支持可配置的噪音过滤规则（含中文网站常见广告）
036 */
037public class HtmlExtractor implements FileExtractor {
038
039    private static final Set<String> SUPPORTED_MIME_TYPES;
040    private static final Set<String> SUPPORTED_EXTENSIONS;
041
042    static {
043        Set<String> mimeTypes = new HashSet<>();
044        mimeTypes.add("text/html");
045        mimeTypes.add("application/xhtml+xml");
046        SUPPORTED_MIME_TYPES = Collections.unmodifiableSet(mimeTypes);
047
048        Set<String> extensions = new HashSet<>();
049        extensions.add("html");
050        extensions.add("htm");
051        extensions.add("xhtml");
052        extensions.add("mhtml");
053        SUPPORTED_EXTENSIONS = Collections.unmodifiableSet(extensions);
054    }
055
056    // 噪音过滤规则
057    private static final Set<String> DEFAULT_SELECTORS = Collections.unmodifiableSet(new HashSet<>(Arrays.asList(
058        "script", "style", "noscript",
059        "nav", "header", "footer", "aside",
060        "iframe", "embed", "object", "video", "audio",
061        ".ads", ".advertisement", ".ad-", "ad:",
062        ".sidebar", ".sider", ".widget", ".module",
063        ".breadcrumb", ".pager", ".pagination",
064        ".share", ".social", ".like", ".subscribe",
065        ".cookie", ".consent", ".banner", ".popup",
066        "[data-ad]", "[data-testid*='ad']", "[data-type='advertisement']"
067    )));
068
069    private static final Set<String> CLASS_KEYWORDS = Collections.unmodifiableSet(new HashSet<>(Arrays.asList(
070        "ad", "adv", "advertisement", "banner", "sponsor",
071        "sidebar", "sider", "widget", "module", "recommend",
072        "related", "similar", "youlike", "hot", "tuijian",
073        "share", "social", "like", "follow", "subscribe",
074        "cookie", "consent", "popup", "modal", "dialog",
075        "footer", "nav", "breadcrumb", "pager", "pagination"
076    )));
077
078    private static final Pattern ID_CLASS_PATTERN = Pattern.compile("(?i)\\b(" +
079        String.join("|",
080            "ad", "adv", "advertisement", "banner", "sponsor",
081            "sidebar", "sider", "widget", "module", "tuijian",
082            "share", "social", "like", "follow", "subscribe",
083            "cookie", "consent", "popup", "modal", "dialog",
084            "footer", "nav", "breadcrumb", "pager", "pagination"
085        ) + ")\\b"
086    );
087
088    // 可动态添加的自定义规则
089    private static final Set<String> CUSTOM_SELECTORS = ConcurrentHashMap.newKeySet();
090    private static final Set<String> CUSTOM_CLASS_KEYWORDS = ConcurrentHashMap.newKeySet();
091
092    /**
093     * 添加自定义噪音选择器（CSS 选择器）
094     */
095    public static void addCustomSelector(String selector) {
096        CUSTOM_SELECTORS.add(selector);
097    }
098
099    /**
100     * 添加自定义 class/id 关键词
101     */
102    public static void addCustomKeyword(String keyword) {
103        CUSTOM_CLASS_KEYWORDS.add(keyword.toLowerCase());
104    }
105
106    @Override
107    public boolean supports(DocumentSource source) {
108        String mimeType = source.getMimeType();
109        String fileName = source.getFileName();
110
111        if (mimeType != null && SUPPORTED_MIME_TYPES.contains(mimeType.toLowerCase())) {
112            return true;
113        }
114
115        if (fileName != null) {
116            String ext = getExtension(fileName);
117            if (ext != null && SUPPORTED_EXTENSIONS.contains(ext.toLowerCase())) {
118                return true;
119            }
120        }
121
122        return false;
123    }
124
125    @Override
126    public String extractText(DocumentSource source) throws IOException {
127        try (InputStream is = source.openStream()) {
128            String html = EncodingDetectUtil.readToString(is);
129            if (html.trim().isEmpty()) {
130                return "";
131            }
132
133            Document doc = Jsoup.parse(html);
134            doc.outputSettings().prettyPrint(false);
135
136            StringBuilder text = new StringBuilder();
137
138            extractTitle(doc, text);
139            extractBodyContent(doc, text);
140
141            return text.toString().trim();
142        } catch (Exception e) {
143            throw new IOException("Failed to parse HTML: " + e.getMessage(), e);
144        }
145    }
146
147    private void extractTitle(Document doc, StringBuilder text) {
148        Elements titleEl = doc.select("title");
149        if (!titleEl.isEmpty()) {
150            String title = titleEl.first().text().trim();
151            if (!title.isEmpty()) {
152                text.append(title).append("\n\n");
153            }
154        }
155    }
156
157    private void extractBodyContent(Document doc, StringBuilder text) {
158        Element body = doc.body();
159        if (body == null) return;
160
161        // 1. 移除已知噪音元素（CSS 选择器）
162        removeElementsBySelectors(body);
163
164        // 2. 移除 class/id 包含关键词的元素
165        removeElementsWithKeywords(body);
166
167        // 3. 遍历剩余节点
168        for (Node node : body.childNodes()) {
169            appendNodeText(node, text, 0);
170        }
171    }
172
173    /**
174     * 使用 CSS 选择器移除噪音
175     */
176    private void removeElementsBySelectors(Element body) {
177        List<String> allSelectors = new ArrayList<>(DEFAULT_SELECTORS);
178        allSelectors.addAll(CUSTOM_SELECTORS);
179
180        for (String selector : allSelectors) {
181            try {
182                body.select(selector).remove();
183            } catch (Exception e) {
184                // 忽略无效选择器
185            }
186        }
187    }
188
189    /**
190     * 移除 class 或 id 包含关键词的元素
191     */
192    private void removeElementsWithKeywords(Element body) {
193        // 合并默认和自定义关键词
194        Set<String> keywords = new HashSet<>(CLASS_KEYWORDS);
195        keywords.addAll(CUSTOM_CLASS_KEYWORDS);
196
197        // 使用 DFS 遍历所有元素
198        Deque<Element> stack = new ArrayDeque<>();
199        stack.push(body);
200
201        while (!stack.isEmpty()) {
202            Element el = stack.pop();
203
204            // 检查 class 或 id 是否匹配
205            String classes = el.className().toLowerCase();
206            String id = el.id().toLowerCase();
207
208            for (String keyword : keywords) {
209                if (classes.contains(keyword) || id.contains(keyword)) {
210                    el.remove(); // 移除整个元素
211                    break;
212                }
213            }
214
215            // 匹配正则模式
216            if (ID_CLASS_PATTERN.matcher(classes).find() ||
217                ID_CLASS_PATTERN.matcher(id).find()) {
218                el.remove();
219                continue;
220            }
221
222            // 将子元素加入栈
223            for (Element child : el.children()) {
224                stack.push(child);
225            }
226        }
227    }
228
229
230    private String repeat(String string, int times) {
231        if (times <= 0) return "";
232        StringBuilder sb = new StringBuilder();
233        for (int i = 0; i < times; i++) {
234            sb.append(string);
235        }
236        return sb.toString();
237    }
238
239    // 节点文本提取
240    private void appendNodeText(Node node, StringBuilder text, int level) {
241        if (node == null) return;
242
243        if (node instanceof org.jsoup.nodes.TextNode) {
244            String txt = ((org.jsoup.nodes.TextNode) node).text().trim();
245            if (!txt.isEmpty()) {
246                text.append(txt).append(" ");
247            }
248        } else if (node instanceof Element) {
249            Element el = (Element) node;
250            String tagName = el.tagName().toLowerCase();
251
252            if (tagName.matches("h[1-6]")) {
253                text.append("\n")
254                    .append(repeat("##", Integer.parseInt(tagName.substring(1))))
255                    .append(el.text().trim())
256                    .append("\n\n");
257            } else if (tagName.equals("p")) {
258                String paraText = el.text().trim();
259                if (!paraText.isEmpty()) {
260                    text.append(paraText).append("\n\n");
261                }
262            } else if (tagName.equals("li")) {
263                text.append("- ").append(el.text().trim()).append("\n");
264            } else if (tagName.equals("table")) {
265                extractTable(el, text);
266                text.append("\n");
267            } else if (tagName.equals("br")) {
268                text.append("\n");
269            } else if (tagName.equals("a")) {
270                String href = el.attr("href");
271                String textPart = el.text().trim();
272                if (!textPart.isEmpty()) {
273                    text.append(textPart);
274                    if (!href.isEmpty() && !href.equals(textPart)) {
275                        text.append(" [").append(href).append("]");
276                    }
277                    text.append(" ");
278                }
279            } else if (isBlockLevel(tagName)) {
280                text.append("\n");
281                for (Node child : el.childNodes()) {
282                    appendNodeText(child, text, level + 1);
283                }
284                text.append("\n");
285            } else {
286                for (Node child : el.childNodes()) {
287                    appendNodeText(child, text, level);
288                }
289            }
290        }
291    }
292
293    private boolean isBlockLevel(String tagName) {
294        Set<String> blockTags = new HashSet<>(Arrays.asList(
295            "div", "p", "h1", "h2", "h3", "h4", "h5", "h6",
296            "ul", "ol", "li", "table", "tr", "td", "th",
297            "blockquote", "pre", "section", "article", "figure"
298        ));
299        return blockTags.contains(tagName);
300    }
301
302    private void extractTable(Element table, StringBuilder text) {
303        text.append("[Table Start]\n");
304        Elements rows = table.select("tr");
305        for (Element row : rows) {
306            Elements cells = row.select("td, th");
307            List<String> cellTexts = new ArrayList<>();
308            for (Element cell : cells) {
309                cellTexts.add(cell.text().trim());
310            }
311            text.append(String.join(" | ", cellTexts)).append("\n");
312        }
313        text.append("[Table End]\n");
314    }
315
316
317    @Override
318    public int getOrder() {
319        return 12;
320    }
321
322    private String getExtension(String fileName) {
323        if (fileName == null || !fileName.contains(".")) return null;
324        int lastDot = fileName.lastIndexOf('.');
325        return fileName.substring(lastDot + 1);
326    }
327}