package cn.iocoder.yudao.module.qcreport.service.aiimport.document;

import cn.iocoder.yudao.module.qcreport.config.QcReportAiImportProperties;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.apache.pdfbox.Loader;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.rendering.ImageType;
import org.apache.pdfbox.rendering.PDFRenderer;
import org.apache.pdfbox.text.PDFTextStripper;
import org.springframework.core.annotation.Order;
import org.springframework.stereotype.Component;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

import static cn.iocoder.yudao.framework.common.exception.util.ServiceExceptionUtil.exception;
import static cn.iocoder.yudao.module.qcreport.enums.ErrorCodeConstants.AI_IMPORT_FILE_EMPTY;
import static cn.iocoder.yudao.module.qcreport.enums.ErrorCodeConstants.AI_IMPORT_TOO_MANY_PAGES;

/**
 * PDF 抽取器，双通道。
 * <p>
 * <b>通道判定依据是「这份 PDF 有没有文本层」，不是「它是不是 PDF」。</b>
 * 这个区别很实际：电子版 PDF 抽出文本只要几十 KB、读得也准；扫描版 PDF 抽出来是空的，
 * 只能把页面渲染成图交给多模态模型——成本高一个数量级，但这是唯一能读出字的办法。
 * 判据是平均每页字符数超过阈值。
 * <p>
 * 刻意<b>不</b>做混合型 PDF 的逐页通路（有的页有文本层、有的没有）：那需要按页决定通道，
 * 调用次数翻倍而收益只在少数文件上。整份文件按一个通道走，行为可预期。
 */
@Slf4j
@Component
@Order(2)
public class PdfImportAdapter implements QcReportImportAdapter {

    @Resource
    private QcReportAiImportProperties properties;

    @Override
    public boolean supports(String extension, String contentType) {
        return "pdf".equals(extension) || "application/pdf".equalsIgnoreCase(contentType);
    }

    @Override
    public QcReportDocumentExtract extract(byte[] content, String originalFilename) {
        if (content == null || content.length == 0) {
            throw exception(AI_IMPORT_FILE_EMPTY, originalFilename, 0);
        }
        try (PDDocument document = Loader.loadPDF(content)) {
            int pageCount = document.getNumberOfPages();
            if (pageCount <= 0) {
                throw exception(AI_IMPORT_FILE_EMPTY, originalFilename, 0);
            }
            // 超限直接报错而不是截断：静默截断会让用户以为整份文件都识别过了，
            // 拿到的模板少了几页却没有任何提示，比报错危险得多
            if (pageCount > properties.getMaxPagesPerFile()) {
                throw exception(AI_IMPORT_TOO_MANY_PAGES, originalFilename, pageCount,
                        properties.getMaxPagesPerFile());
            }
            List<String> pageTexts = extractPageTexts(document, pageCount);
            int totalChars = pageTexts.stream().mapToInt(t -> t.replaceAll("\\s", "").length()).sum();
            int threshold = properties.getPdfTextPageThreshold() * pageCount;
            if (totalChars >= threshold) {
                return QcReportDocumentExtract.ofText(joinPages(pageTexts), pageCount, "pdfbox-text", List.of());
            }
            // 走图片通道：把「为什么慢/为什么贵」讲清楚，用户才知道换电子版 PDF 能快很多
            return renderPages(document, pageCount, originalFilename,
                    List.of("文件「" + originalFilename + "」没有文本层（共 " + pageCount
                            + " 页），已按扫描件逐页识别，会比较慢"));
        } catch (IOException e) {
            log.warn("PDF 解析失败，file={}", originalFilename, e);
            throw exception(AI_IMPORT_FILE_EMPTY, originalFilename, 0);
        }
    }

    /**
     * 逐页抽文本。用 {@code setSortByPosition(true)} 让阅读顺序按版面坐标排，
     * 否则表格里的文字会按 PDF 内部绘制顺序乱序抽出，模型看到的列就错位了。
     */
    private List<String> extractPageTexts(PDDocument document, int pageCount) throws IOException {
        List<String> pageTexts = new ArrayList<>(pageCount);
        for (int i = 0; i < pageCount; i++) {
            PDFTextStripper stripper = new PDFTextStripper();
            stripper.setSortByPosition(true);
            stripper.setStartPage(i + 1);
            stripper.setEndPage(i + 1);
            pageTexts.add(stripper.getText(document).trim());
        }
        return pageTexts;
    }

    private String joinPages(List<String> pageTexts) {
        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < pageTexts.size(); i++) {
            if (i > 0) {
                sb.append("\n--- 第 ").append(i + 1).append(" 页 ---\n");
            }
            sb.append(pageTexts.get(i));
        }
        return sb.toString();
    }

    /**
     * 逐页渲染成 PNG。
     * <p>
     * 渲染失败不吞：回落去试文本通道，两条路都不通才报「未解析出任何可用内容」。
     * 之所以要回落，是因为「字符数低于阈值」可能只是文件用的字体没带 ToUnicode 表
     * （抽出来是乱码或空），但渲染本身还是好的——不过那种情况下回落取到的文本也是空的，
     * 所以真正会落到 {@link #extractPageTexts} 的场景是渲染器对这个文件不工作。
     */
    private QcReportDocumentExtract renderPages(PDDocument document, int pageCount,
                                                String originalFilename, List<String> notes) {
        List<QcReportDocumentExtract.ImagePart> images = new ArrayList<>(pageCount);
        try {
            PDFRenderer renderer = new PDFRenderer(document);
            for (int i = 0; i < pageCount; i++) {
                BufferedImage image = renderer.renderImageWithDPI(i, properties.getRenderDpi(), ImageType.RGB);
                ByteArrayOutputStream out = new ByteArrayOutputStream();
                ImageIO.write(image, "png", out);
                images.add(new QcReportDocumentExtract.ImagePart(
                        out.toByteArray(), "image/png", "第 " + (i + 1) + " 页"));
            }
            return QcReportDocumentExtract.ofImages(images, pageCount, "pdfbox-render", notes);
        } catch (IOException | RuntimeException e) {
            log.warn("PDF 渲染扫描页失败，尝试回落文本通道，file={}", originalFilename, e);
            try {
                List<String> pageTexts = extractPageTexts(document, pageCount);
                if (pageTexts.stream().anyMatch(t -> !t.isBlank())) {
                    List<String> fallbackNotes = new ArrayList<>(notes);
                    fallbackNotes.add("文件「" + originalFilename + "」的页面无法渲染成图片，已改为读取其文本层");
                    return QcReportDocumentExtract.ofText(joinPages(pageTexts), pageCount,
                            "pdfbox-text-fallback", fallbackNotes);
                }
            } catch (IOException ignored) {
                // 回落也读不出来，走下面的统一报错
            }
            throw exception(AI_IMPORT_FILE_EMPTY, originalFilename, pageCount);
        }
    }

}