package com.zhyc.module.app.util;

import ai.djl.modality.cv.Image;
import cn.smartjavaai.common.cv.SmartImageFactory;
import cn.smartjavaai.common.enums.DeviceEnum;
import cn.smartjavaai.common.utils.ImageUtils;
import cn.smartjavaai.ocr.config.DirectionModelConfig;
import cn.smartjavaai.ocr.config.OcrDetModelConfig;
import cn.smartjavaai.ocr.config.OcrRecModelConfig;
import cn.smartjavaai.ocr.config.OcrRecOptions;
import cn.smartjavaai.ocr.entity.OcrInfo;
import cn.smartjavaai.ocr.enums.CommonDetModelEnum;
import cn.smartjavaai.ocr.enums.CommonRecModelEnum;
import cn.smartjavaai.ocr.enums.DirectionModelEnum;
import cn.smartjavaai.ocr.factory.OcrModelFactory;
import cn.smartjavaai.ocr.model.common.detect.OcrCommonDetModel;
import cn.smartjavaai.ocr.model.common.direction.OcrDirectionModel;
import cn.smartjavaai.ocr.model.common.recognize.OcrCommonRecModel;
import com.alibaba.fastjson.JSONObject;
import com.zhyc.common.config.RuoYiConfig;
import lombok.extern.slf4j.Slf4j;
import org.junit.BeforeClass;
import org.junit.Test;

import java.io.File;
import java.io.IOException;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.List;

/**
 * OCR 文本识别 示例
 * 模型下载地址：https://pan.baidu.com/s/1MLfd73Vjdpnuls9-oqc9uw?pwd=1234 提取码: 1234
 * 开发文档：http://doc.smartjavaai.cn/
 */
/**
 * OCR识别工具类
 * 提供文本识别相关的功能，包括获取不同类型的识别模型、文本检测模型、方向检测模型等
 * 支持普通识别、手写识别、带方向矫正的识别以及批量识别等功能
 */
@Slf4j
public class OcrRecognizeUtil {

    //设备类型
    public static DeviceEnum device =DeviceEnum.CPU;


    /**
     * 测试初始化方法，在所有测试用例执行前调用一次
     * 设置SmartImageFactory使用的引擎为OpenCV，并可选配置缓存路径
     *
     * @throws IOException 当初始化过程中发生IO异常时抛出
     */
    @BeforeClass
    public static void beforeAll() throws IOException {
        SmartImageFactory.setEngine(SmartImageFactory.Engine.OPENCV);
        //修改缓存路径
        //Config.setCachePath("/Users/xxx/smartjavaai_cache");
    }

    /**
     * 获取通用识别模型（高精确度模型）
     * 注意事项：高精度模型，识别准确度高，速度慢
     *
     * @return 返回高精度的OCR通用识别模型实例
     */
    public OcrCommonRecModel getProRecModel(){
        OcrRecModelConfig recModelConfig = new OcrRecModelConfig();
        //指定文本识别模型，切换模型需要同时修改modelEnum及modelPath
        recModelConfig.setRecModelEnum(CommonRecModelEnum.PP_OCR_V5_SERVER_REC_MODEL);
        //指定识别模型位置，需要更改为自己的模型路径（下载地址请查看文档）
        recModelConfig.setRecModelPath("C:/wwwroot/ocr/rec/PP-OCRv5_server_rec_infer/PP-OCRv5_server_rec.onnx");
        recModelConfig.setDevice(device);
        recModelConfig.setTextDetModel(getProDetectionModel());
        recModelConfig.setDirectionModel(getDirectionModel());
        return OcrModelFactory.getInstance().getRecModel(recModelConfig);
    }

    /**
     * 获取通用识别模型（极速模型）
     * 注意事项：极速模型，识别准确度低，速度快
     *
     * @return 返回极速的OCR通用识别模型实例
     */
    public OcrCommonRecModel getFastRecModel(){
        OcrRecModelConfig recModelConfig = new OcrRecModelConfig();
        //指定文本识别模型，切换模型需要同时修改modelEnum及modelPath
        recModelConfig.setRecModelEnum(CommonRecModelEnum.PP_OCR_V5_MOBILE_REC_MODEL);
        //指定识别模型位置，需要更改为自己的模型路径（下载地址请查看文档）
        recModelConfig.setRecModelPath("C:/wwwroot/ocr/rec/PP-OCRv5_mobile_rec_infer/PP-OCRv5_mobile_rec_infer.onnx");
        recModelConfig.setDevice(device);
        recModelConfig.setTextDetModel(getFastDetectionModel());
        recModelConfig.setDirectionModel(getDirectionModel()); // 添加方向模型配置
        return OcrModelFactory.getInstance().getRecModel(recModelConfig);
    }


    /**
     * 获取文本检测模型(极速模型)
     * 注意事项：极速模型，识别准确度低，速度快
     *
     * @return 返回极速的OCR通用检测模型实例
     */
    public OcrCommonDetModel getFastDetectionModel() {
        OcrDetModelConfig config = new OcrDetModelConfig();
        //指定检测模型，切换模型需要同时修改modelEnum及modelPath
        config.setModelEnum(CommonDetModelEnum.PP_OCR_V5_MOBILE_DET_MODEL);
        //指定模型位置，需要更改为自己的模型路径（下载地址请查看文档）
        config.setDetModelPath("C:/wwwroot/ocr/det/PP-OCRv5_mobile_det_infer/PP-OCRv5_mobile_det_infer.onnx");
        config.setDevice(device);
        return OcrModelFactory.getInstance().getDetModel(config);
    }

    /**
     * 获取文本检测模型(高精确度模型)
     * 注意事项：高精度模型，识别准确度高，速度慢
     *
     * @return 返回高精度的OCR通用检测模型实例
     */
    public OcrCommonDetModel getProDetectionModel() {
        OcrDetModelConfig config = new OcrDetModelConfig();
        //指定检测模型，切换模型需要同时修改modelEnum及modelPath
        config.setModelEnum(CommonDetModelEnum.PP_OCR_V5_SERVER_DET_MODEL);
        //指定模型位置，需要更改为自己的模型路径（下载地址请查看文档）
        config.setDetModelPath("C:/wwwroot/ocr/det/PP-OCRv5_server_det_infer/PP-OCRv5_server_det.onnx");
        config.setDevice(device);
        return OcrModelFactory.getInstance().getDetModel(config);
    }

    /**
     * 获取方向检测模型
     *
     * @return 返回OCR方向检测模型实例
     */
    public OcrDirectionModel getDirectionModel(){
        DirectionModelConfig directionModelConfig = new DirectionModelConfig();
        //指定行文本方向检测模型，切换模型需要同时修改modelEnum及modelPath
        directionModelConfig.setModelEnum(DirectionModelEnum.PP_LCNET_X0_25);
        //指定行文本方向检测模型路径，需要更改为自己的模型路径（下载地址请查看文档）
        directionModelConfig.setModelPath("C:/wwwroot/ocr/ori/PP-LCNet_x0_25_textline_ori_infer.onnx");
        directionModelConfig.setDevice(device);
        return OcrModelFactory.getInstance().getDirectionModel(directionModelConfig);
    }



    /**
     * 文本识别
     * 支持简体中文、繁体中文、英文、日文四种主要语言，以及手写、竖版、拼音、生僻字
     * 流程：文本检测 -> 文本识别
     * 注意事项：
     * 1、批量检测时，模型应统一放在外层 try 中使用，避免重复加载，自动释放资源更安全。
     * 2、模型文件需要放在单独文件夹
     */
    public String recognize(String path){
        try {
            OcrCommonRecModel recModel = getFastRecModel();
            //不带方向矫正，分行返回文本
            OcrRecOptions options = new OcrRecOptions(false, true);
            //创建Image对象，可以从文件、url、InputStream创建、BufferedImage、Base64创建，具体使用方法可以查看文档
            Image image = SmartImageFactory.getInstance().fromFile(path);
            OcrInfo ocrInfo = recModel.recognize(image, options);
            log.info("OCR识别结果：{}", JSONObject.toJSONString(ocrInfo.getFullText()));

            // 识别成功后删除原图片
            deleteFile(path);
            return ocrInfo.getFullText();
        } catch (Exception e) {
            e.printStackTrace();
        }
       return null;
    }


    /**
     * 文本识别（手写字）
     * 支持简体中文、繁体中文、英文、日文四种主要语言，以及手写、竖版、拼音、生僻字
     * 流程：文本检测 -> 文本识别
     * 注意事项：
     * 1、批量检测时，模型应统一放在外层 try 中使用，避免重复加载，自动释放资源更安全。
     * 2、模型文件需要放在单独文件夹
     * 3、识别完成后会删除原图片文件
     */
    public void recognizeHandWriting(String path){
        try {
            OcrCommonRecModel recModel = getFastRecModel();
            //创建Image对象，可以从文件、url、InputStream创建、BufferedImage、Base64创建，具体使用方法可以查看文档
            Image image = SmartImageFactory.getInstance().fromFile(path);
            OcrInfo ocrInfo = recModel.recognize(image, new OcrRecOptions());
            log.info("OCR识别结果：{}", JSONObject.toJSONString(ocrInfo));

            // 识别成功后删除原图片
            deleteFile(path);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

/**
 * 文本识别（带方向矫正）
 * 支持简体中文、繁体中文、英文、日文四种主要语言，以及手写、竖版、拼音、生僻字
 * 本方法支持多角度文字识别
 * 流程：文本检测 -> 方向检测 -> 方向矫正 ->  文本识别
 * 注意事项：
 * 1、批量检测时，模型应统一放在外层 try 中使用，避免重复加载，自动释放资源更安全。
 * 2、模型文件需要放在单独文件夹
 * 3、识别完成后会删除原图片文件
 */
public String recognizeTwo(String path){
    try {
        OcrCommonRecModel recModel = getFastRecModel();
        //带方向矫正，分行返回文本
        OcrRecOptions options = new OcrRecOptions(true, true);
        //创建Image对象，可以从文件、url、InputStream创建、BufferedImage、Base64创建，具体使用方法可以查看文档
        Image image = SmartImageFactory.getInstance().fromFile(path);
        OcrInfo ocrInfo = recModel.recognize(image, options);
        log.info("OCR识别结果：{}", JSONObject.toJSONString(ocrInfo));

        // 识别成功后删除原图片
        deleteFile(path);
        return ocrInfo.getFullText();
    } catch (Exception e) {
        e.printStackTrace();
    }
    return null;
}


    /**
     * 删除指定路径的文件
     *
     * @param filePath 要删除的文件路径
     * @return 删除是否成功
     */
    private boolean deleteFile(String filePath) {
        try {
            File file = new File(filePath);
            if (file.exists()) {
                boolean deleted = file.delete();
                if (deleted) {
                    log.info("已删除文件: {}", filePath);
                } else {
                    log.warn("删除文件失败: {}", filePath);
                }
                return deleted;
            } else {
                log.warn("文件不存在: {}", filePath);
                return false;
            }
        } catch (Exception e) {
            log.error("删除文件时发生异常: {}", filePath, e);
            return false;
        }
    }


}
