package com.xly.ocr.test;

import com.benjaminwan.ocrlibrary.OcrResult;
import com.benjaminwan.ocrlibrary.Point;
import com.benjaminwan.ocrlibrary.TextBlock;
import io.github.mymonstercat.Model;
import io.github.mymonstercat.ocr.InferenceEngine;
import io.github.mymonstercat.ocr.config.ParamConfig;

import javax.imageio.ImageIO;
import java.awt.*;
import java.awt.image.BufferedImage;
import java.awt.image.Kernel;
import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

public class Main {

    static {
        try {
            String customTempDir = "D:/temp/ocrJava";
            File tempDir = new File(customTempDir);
            if (!tempDir.exists()) {
                tempDir.mkdirs();
            }
            System.setProperty("java.io.tmpdir", customTempDir);
            System.setProperty("TMP", customTempDir);
            System.setProperty("TEMP", customTempDir);
            System.setProperty("ORT_TMP_DIR", customTempDir);

            System.out.println("==================================");
            System.out.println("java.io.tmpdir: " + System.getProperty("java.io.tmpdir"));
            System.out.println("临时目录是否存在: " + tempDir.exists());
            System.out.println("==================================");

        } catch (Exception e) {
            System.err.println("设置临时目录失败: " + e.getMessage());
            e.printStackTrace();
        }
    }

    public static void main(String[] args) {
        try {
            System.out.println("OCR 程序开始执行...");

            // 检查旧的临时目录
            String oldTempPath = "C:\\Users\\钱豹\\AppData\\Local\\Temp\\ocrJava";
            File oldTempDir = new File(oldTempPath);
            if (oldTempDir.exists()) {
                System.out.println("发现旧的临时目录: " + oldTempPath);
            }

            // 1. 初始化引擎（V4 模型）
            System.out.println("正在初始化 OCR 引擎...");
            InferenceEngine engine = InferenceEngine.getInstance(Model.ONNX_PPOCR_V4);

            // 2. 创建优化的参数配置
            ParamConfig config = createOptimizedParamConfig();

            // 3. 设置图片路径
            String imagePath = "E:/aa/b.jpg";

            File imageFile = new File(imagePath);
            if (!imageFile.exists()) {
                System.err.println("图片文件不存在: " + imagePath);
                return;
            }

            // 4. 图像预处理（可选，注释掉可提高速度）
            System.out.println("正在进行图像预处理...");
            String processedImagePath = preprocessImage(imagePath);

            // 5. 执行识别
            System.out.println("开始识别图片: " + processedImagePath);
            long startTime = System.currentTimeMillis();
            OcrResult ocrResult = engine.runOcr(processedImagePath, config);
            long endTime = System.currentTimeMillis();

            // 6. 输出识别结果
            String text = ocrResult.getStrRes().trim();
            System.out.println("==================================");
            System.out.println("识别结果：");
            System.out.println(text);
            System.out.println("==================================");
            System.out.println("识别耗时: " + (endTime - startTime) + " ms");

            // 7. 输出文本块详细信息（可选，用于调试）
            if (ocrResult.getTextBlocks() != null && !ocrResult.getTextBlocks().isEmpty()) {
                System.out.println("\n文本块详情（共" + ocrResult.getTextBlocks().size() + "块）：");
                List<TextBlock> textBlocks = ocrResult.getTextBlocks();
                for (int i = 0; i < textBlocks.size(); i++) {
                    TextBlock block = textBlocks.get(i);
                    System.out.printf("  块%d: %s (置信度: %.2f)%n",
                            i + 1,
                            block.getText(),
                            block.getBoxScore()
                    );
                }
            }

            // 8. 清理临时文件
            if (!processedImagePath.equals(imagePath)) {
                File processedFile = new File(processedImagePath);
                if (processedFile.exists()) {
                    processedFile.delete();
                    System.out.println("\n已清理临时文件: " + processedImagePath);
                }
            }

        } catch (Exception e) {
            System.err.println("OCR 识别失败: " + e.getMessage());
            e.printStackTrace();
        }
    }

    private static ParamConfig createOptimizedParamConfig() {
        ParamConfig config = new ParamConfig();

        // 手写体专用超强参数
        config.setPadding(100);
        config.setMaxSideLen(1200);

        // 极低阈值 = 不漏检任何手写文字
        config.setBoxScoreThresh(0.25f);
        config.setBoxThresh(0.15f);
        config.setUnClipRatio(2.5f);

        // 必须开启角度矫正
        config.setDoAngle(true);
        config.setMostAngle(true);

        return config;
    }

    /**
     * 图像预处理：提高OCR识别准确度
     */
    private static String preprocessImage(String imagePath) throws IOException {
        File inputFile = new File(imagePath);
        BufferedImage originalImage = ImageIO.read(inputFile);

        if (originalImage == null) return imagePath;

        // 固定缩放到最佳尺寸
        BufferedImage resized = resizeImageWithQuality(originalImage, 1000, 1500);

        // 手写体必须：二值化（黑白强化）+ 锐化
        BufferedImage binary = toBinary(resized);  // 黑白强化，核心！
        BufferedImage sharp = sharpenImage(binary); // 锐化

        // 保存
        String outputPath = imagePath.substring(0, imagePath.lastIndexOf(".")) + "_final.jpg";
        ImageIO.write(sharp, "jpg", new File(outputPath));
        return outputPath;
    }

    // 手写体专用：黑白二值化，文字瞬间清晰
    private static BufferedImage toBinary(BufferedImage image) {
        BufferedImage gray = new BufferedImage(image.getWidth(), image.getHeight(), BufferedImage.TYPE_BYTE_GRAY);
        Graphics g = gray.getGraphics();
        g.drawImage(image, 0, 0, null);
        g.dispose();

        // 手写体阈值 120，文字最清晰
        int threshold = 120;
        for (int y = 0; y < gray.getHeight(); y++) {
            for (int x = 0; x < gray.getWidth(); x++) {
                int rgb = gray.getRGB(x, y);
                int grayValue = (rgb >> 16) & 0xFF;
                if (grayValue < threshold) {
                    gray.setRGB(x, y, Color.BLACK.getRGB());
                } else {
                    gray.setRGB(x, y, Color.WHITE.getRGB());
                }
            }
        }
        return gray;
    }

    /**
     * 高质量调整图片大小
     */
    private static BufferedImage resizeImageWithQuality(BufferedImage original, int maxWidth, int maxHeight) {
        int width = original.getWidth();
        int height = original.getHeight();

        // 如果图片尺寸合适，不进行调整
        if (width <= maxWidth && height <= maxHeight) {
            return original;
        }

        // 计算缩放比例
        double ratio = Math.min((double) maxWidth / width, (double) maxHeight / height);
        int newWidth = (int) (width * ratio);
        int newHeight = (int) (height * ratio);

        // 创建缩放后的图片
        BufferedImage resized = new BufferedImage(newWidth, newHeight, BufferedImage.TYPE_INT_RGB);
        Graphics2D g = resized.createGraphics();
        g.setRenderingHint(RenderingHints.KEY_INTERPOLATION, RenderingHints.VALUE_INTERPOLATION_BICUBIC);
        g.setRenderingHint(RenderingHints.KEY_RENDERING, RenderingHints.VALUE_RENDER_QUALITY);
        g.setRenderingHint(RenderingHints.KEY_ANTIALIASING, RenderingHints.VALUE_ANTIALIAS_ON);
        g.drawImage(original, 0, 0, newWidth, newHeight, null);
        g.dispose();

        System.out.println("图片已缩放: " + width + "x" + height + " -> " + newWidth + "x" + newHeight);
        return resized;
    }

    /**
     * 增强对比度
     */
    private static BufferedImage enhanceContrast(BufferedImage image) {
        BufferedImage result = new BufferedImage(image.getWidth(), image.getHeight(), image.getType());

        for (int y = 0; y < image.getHeight(); y++) {
            for (int x = 0; x < image.getWidth(); x++) {
                Color color = new Color(image.getRGB(x, y));

                // 增强对比度
                int red = (int) (color.getRed() * 1.2);
                int green = (int) (color.getGreen() * 1.2);
                int blue = (int) (color.getBlue() * 1.2);

                // 限制RGB值范围
                red = Math.min(255, Math.max(0, red));
                green = Math.min(255, Math.max(0, green));
                blue = Math.min(255, Math.max(0, blue));

                result.setRGB(x, y, new Color(red, green, blue).getRGB());
            }
        }

        System.out.println("对比度已增强");
        return result;
    }

    /**
     * 锐化图像（可选，使文字边缘更清晰）
     */
    private static BufferedImage sharpenImage(BufferedImage image) {
        BufferedImage result = new BufferedImage(image.getWidth(), image.getHeight(), image.getType());

        float[] sharpenKernel = {
                0, -1, 0,
                -1, 5, -1,
                0, -1, 0
        };

        Kernel kernel = new Kernel(3, 3, sharpenKernel);
        java.awt.image.ConvolveOp op = new java.awt.image.ConvolveOp(kernel, java.awt.image.ConvolveOp.EDGE_NO_OP, null);
        op.filter(image, result);

        System.out.println("图像已锐化");
        return result;
    }
}