package org.dromara.system.ai.loader;

import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.apache.poi.hwpf.HWPFDocument;
import org.apache.poi.hwpf.extractor.WordExtractor;
import org.apache.poi.xwpf.extractor.XWPFWordExtractor;
import org.apache.poi.xwpf.usermodel.XWPFDocument;
import org.dromara.system.ai.splitter.TextSplitter;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.io.InputStream;
import java.util.List;

@Component
@AllArgsConstructor
@Slf4j
public class WordLoader implements ResourceLoader {

    private final TextSplitter textSplitter;

    @Override
    public String getContent(InputStream inputStream) {
        try {
            if (isDocFile(inputStream)) {
                // 处理 .doc 文件
                inputStream.reset();
                HWPFDocument document = new HWPFDocument(inputStream);
                WordExtractor wordExtractor = new WordExtractor(document);
                String content = wordExtractor.getText();
                wordExtractor.close();
                return content;
            } else {
                // 处理 .docx文件
                inputStream.reset();
                XWPFDocument document = new XWPFDocument(inputStream);
                XWPFWordExtractor xwpfWordExtractor = new XWPFWordExtractor(document);
                String content = xwpfWordExtractor.getText();
                xwpfWordExtractor.close();
                return content;
            }
        } catch (IOException e) {
            throw new RuntimeException("解析Word文档失败", e);
        }
    }

    @Override
    public List<String> getChunkList(String content, String kid) {
        return textSplitter.split(content, kid);
    }

    // 判断是否为 .doc 文件的方法
    private boolean isDocFile(InputStream inputStream) throws IOException {
        // 读取文件头部信息判断格式
        byte[] header = new byte[8];
        inputStream.mark(8);
        int bytesRead = inputStream.read(header);
        inputStream.reset();

        if (bytesRead >= 8) {
            // .doc 文件的特征头部
            return header[0] == (byte) 0xD0 && header[1] == (byte) 0xCF &&
                header[2] == (byte) 0x11 && header[3] == (byte) 0xE0;
        }
        return false;
    }
}
