package com.xindao.ocr.swingui.service;

import cn.smartjavaai.common.enums.DeviceEnum;
import com.alibaba.fastjson.JSONObject;
import com.xindao.ocr.smartjavaai.config.DirectionModelConfig;
import com.xindao.ocr.smartjavaai.config.OcrDetModelConfig;
import com.xindao.ocr.smartjavaai.config.OcrRecModelConfig;
import com.xindao.ocr.smartjavaai.config.OcrRecOptions;
import com.xindao.ocr.smartjavaai.entity.OcrInfo;
import com.xindao.ocr.smartjavaai.enums.CommonDetModelEnum;
import com.xindao.ocr.smartjavaai.enums.CommonRecModelEnum;
import com.xindao.ocr.smartjavaai.enums.DirectionModelEnum;
import com.xindao.ocr.smartjavaai.factory.OcrModelFactory;
import com.xindao.ocr.smartjavaai.model.common.detect.OcrCommonDetModel;
import com.xindao.ocr.smartjavaai.model.common.direction.OcrDirectionModel;
import com.xindao.ocr.smartjavaai.model.common.recognize.OcrCommonRecModel;
import com.xindao.ocr.swingui.constant.OcrSwingConstants;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.core.io.ClassPathResource;
import org.springframework.core.io.Resource;
import org.springframework.core.io.support.PathMatchingResourcePatternResolver;
import org.springframework.stereotype.Service;

import javax.annotation.PostConstruct;
import java.io.File;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.StandardCopyOption;

@Slf4j
@Service
@RequiredArgsConstructor
public class OcrService {

    public static DeviceEnum device = DeviceEnum.CPU;

    private final static String MAIN_PP_DIR = "PP_OCRv5";

    private OcrCommonRecModel recModel;

    /**
     * 初始化加载ocr模型
     */
    @PostConstruct
    private void init() {
        log.info("复制ocr文件到本地用户缓存...");
        // 复制PP_OCRv5目录到本地缓存
        copyPaddleCppToCache();
        try {
            recModel = getRecModel();
        } catch (IOException e) {
            log.error("加载OCR模型失败: {}", e.getMessage());
            e.printStackTrace();
        }
    }

    /**
     * 复制resources中的Paddle_CPP目录到本地缓存
     */
    private void copyPaddleCppToCache() {
        try {
            // 创建缓存目录 - 使用用户主目录避免权限问题
            File cacheDir = OcrSwingConstants.ocrDir;
            if (!cacheDir.exists()) {
                cacheDir.mkdirs();
            }

            // 获取resources中的Paddle_CPP目录资源
            PathMatchingResourcePatternResolver resolver = new PathMatchingResourcePatternResolver();
            Resource[] resources = resolver.getResources("classpath:"+MAIN_PP_DIR+"/**");

            // 复制所有资源到缓存目录，保持目录结构
            for (Resource resource : resources) {
                String resourcePath = resource.getURL().getPath();
                // 获取相对于Paddle_CPP的路径
                int startIndex = resourcePath.indexOf(MAIN_PP_DIR) + MAIN_PP_DIR.length();
                String relativePath = resourcePath.substring(startIndex);
                // 处理Windows路径分隔符
                relativePath = relativePath.replace('/', File.separatorChar);
                // 创建目标文件
                File destFile = new File(cacheDir, relativePath);
                // 如果是目录，创建目录
                if (resource.isReadable() && resource.contentLength() == 0 && relativePath.endsWith(File.separator)) {
                    if (!destFile.exists()) {
                        destFile.mkdirs();
                        log.info("创建目录: {}", destFile.getAbsolutePath());
                    }
                } else if (resource.isReadable()) {
                    // 确保父目录存在
                    File parentDir = destFile.getParentFile();
                    if (parentDir != null && !parentDir.exists()) {
                        parentDir.mkdirs();
                    }
                    // 复制文件
                    Files.copy(resource.getInputStream(), destFile.toPath(), StandardCopyOption.REPLACE_EXISTING);
                    log.info("复制资源: {} 到 {}", relativePath, destFile.getAbsolutePath());
                }
            }

        } catch (IOException e) {
            log.error("复制PP_OCRv5目录到缓存失败: {}", e.getMessage());
            e.printStackTrace();
        }
    }


    /**
     * 获取 resources 下模型文件的绝对路径
     *
     * @param relativePath 相对于 resources 的路径
     */
    private String getModelPath(String relativePath) throws IOException {
        String localPath = null;
        try{
            new ClassPathResource(relativePath).getFile().getAbsolutePath();
        }catch (IOException e){
            // 如果找不到文件，则尝试从本地缓存目录获取
            localPath = getLocalPath(relativePath);
            File modelFile = new File(localPath);
            if (!modelFile.exists()) {
                throw new IOException("模型文件不存在: " + localPath);
            }
        } finally {
            log.info("OCR模型文件路径: {}", localPath);
        }
        return localPath;
    }

    /**
     * 获取 resources 下模型文件的用户本地路径
     *
     * @param relativePath 相对于 resources 的路径
     */
    private String getLocalPath(String relativePath) {
        return new File(OcrSwingConstants.ocrDir,File.separator + relativePath).getAbsolutePath();
    }

    public OcrCommonRecModel getRecModel() throws IOException {
        OcrRecModelConfig recModelConfig = new OcrRecModelConfig();
        recModelConfig.setRecModelEnum(CommonRecModelEnum.PP_OCR_V5_MOBILE_REC_MODEL);
        recModelConfig.setRecModelPath(
                getModelPath("PP-OCRv5_server_rec_infer/PP-OCRv5_server_rec.onnx")
        );
        recModelConfig.setDevice(device);
        recModelConfig.setTextDetModel(getDetectionModel());
        return OcrModelFactory.getInstance().getRecModel(recModelConfig);
    }

    public OcrCommonDetModel getDetectionModel() throws IOException {
        OcrDetModelConfig config = new OcrDetModelConfig();
        config.setModelEnum(CommonDetModelEnum.PP_OCR_V5_MOBILE_DET_MODEL);
        config.setDetModelPath(
                getModelPath("PP-OCRv5_server_det_infer/PP-OCRv5_server_det.onnx")
        );
        config.setDevice(device);
        return OcrModelFactory.getInstance().getDetModel(config);
    }

    public OcrDirectionModel getDirectionModel() throws IOException {
        DirectionModelConfig directionModelConfig = new DirectionModelConfig();
        directionModelConfig.setModelEnum(DirectionModelEnum.PP_LCNET_X0_25);
        directionModelConfig.setModelPath(
                getModelPath("PP-LCNet_x0_25_textline_ori_infer/PP-LCNet_x0_25_textline_ori_infer.onnx")
        );
        directionModelConfig.setDevice(device);
        return OcrModelFactory.getInstance().getDirectionModel(directionModelConfig);
    }

    public OcrCommonRecModel getRecModelWithDirection() throws IOException {
        OcrRecModelConfig recModelConfig = new OcrRecModelConfig();
        recModelConfig.setRecModelEnum(CommonRecModelEnum.PP_OCR_V5_MOBILE_REC_MODEL);
        recModelConfig.setRecModelPath(
                getModelPath("PP-OCRv5_mobile_rec_infer/PP-OCRv5_mobile_rec_infer.onnx")
        );
        recModelConfig.setDevice(device);
        recModelConfig.setTextDetModel(getDetectionModel());
        recModelConfig.setDirectionModel(getDirectionModel());
        return OcrModelFactory.getInstance().getRecModel(recModelConfig);
    }

    public String ocr(String url) {
        String fullText = null;
        try {
            OcrRecOptions options = new OcrRecOptions(false, true);
            OcrInfo ocrInfo = recModel.recognize(url, options);
            log.info("OCR识别结果：{}", JSONObject.toJSONString(ocrInfo));
            fullText = ocrInfo.getFullText();
        } catch (Exception e) {
            e.printStackTrace();
        }
        return fullText;
    }

}