package com.xindao.ocr.smartjavaai.model.common.recognize;

import ai.djl.MalformedModelException;
import ai.djl.engine.Engine;
import ai.djl.inference.Predictor;
import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.ImageFactory;
import ai.djl.ndarray.NDArray;
import ai.djl.ndarray.NDManager;
import ai.djl.repository.zoo.Criteria;
import ai.djl.repository.zoo.ModelNotFoundException;
import ai.djl.repository.zoo.ModelZoo;
import ai.djl.repository.zoo.ZooModel;
import cn.hutool.core.img.ImgUtil;
import cn.smartjavaai.common.pool.PredictorFactory;
import cn.smartjavaai.common.utils.FileUtils;
import cn.smartjavaai.common.utils.ImageUtils;
import cn.smartjavaai.common.utils.OpenCVUtils;
import com.xindao.ocr.smartjavaai.config.OcrRecModelConfig;
import com.xindao.ocr.smartjavaai.config.OcrRecOptions;
import com.xindao.ocr.smartjavaai.entity.*;
import com.xindao.ocr.smartjavaai.exception.OcrException;
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.criteria.OcrCommonRecCriterialFactory;
import com.xindao.ocr.smartjavaai.utils.OcrUtils;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.pool2.impl.GenericObjectPool;
import org.opencv.core.Mat;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Objects;
import java.util.stream.Collectors;

/**
 * PPOCRV5 识别模型
 *
 * @author dwj
 */
@Slf4j
public class OcrCommonRecModelImpl implements OcrCommonRecModel {

    private GenericObjectPool<Predictor<Image, String>> recPredictorPool;

    private OcrRecModelConfig config;

    private ZooModel<Image, String> recognitionModel;

    private OcrDirectionModel directionModel;

    private OcrCommonDetModel textDetModel;

    @Override
    public void loadModel(OcrRecModelConfig config) {
        if (StringUtils.isBlank(config.getRecModelPath())) {
            throw new OcrException("recModelPath is null");
        }
        this.config = config;
        this.directionModel = config.getDirectionModel();
        this.textDetModel = config.getTextDetModel();
        //初始化 识别Criteria
        Criteria<Image, String> recCriteria = OcrCommonRecCriterialFactory.createCriteria(config);
        try {
            recognitionModel = ModelZoo.loadModel(recCriteria);
            this.recPredictorPool = new GenericObjectPool<>(new PredictorFactory<>(recognitionModel));
            int predictorPoolSize = config.getPredictorPoolSize();
            if (config.getPredictorPoolSize() <= 0) {
                predictorPoolSize = Runtime.getRuntime().availableProcessors(); // 默认等于CPU核心数
            }
            recPredictorPool.setMaxTotal(predictorPoolSize);
            log.debug("当前设备: " + recognitionModel.getNDManager().getDevice());
            log.debug("当前引擎: " + Engine.getInstance().getEngineName());
            log.debug("模型推理器线程池最大数量: " + predictorPoolSize);
        } catch (IOException | ModelNotFoundException | MalformedModelException e) {
            throw new OcrException("识别模型加载失败", e);
        }

    }


    @Override
    public OcrInfo recognize(String imagePath, OcrRecOptions options) {
        if (StringUtils.isBlank(config.getRecModelPath())) {
            throw new OcrException("recModelPath为空，无法识别");
        }
        if (!FileUtils.isFileExists(imagePath)) {
            throw new OcrException("图像文件不存在");
        }
        Image img = null;
        try {
            img = ImageFactory.getInstance().fromFile(Paths.get(imagePath));
            return recognize(img, options);
        } catch (IOException e) {
            throw new OcrException("无效的图片", e);
        } finally {
            if (img != null) {
                ((Mat) img.getWrappedImage()).release();
            }
        }
    }

    /**
     * @param image
     * @param options
     * @return
     */
    @Override
    public OcrInfo recognize(Image image, OcrRecOptions options) {
        List<OcrInfo> result = batchRecognizeDJLImage(Collections.singletonList(image), options);
        if (CollectionUtils.isEmpty(result)) {
            throw new OcrException("OCR识别结果为空");
        }
        return result.get(0);
    }


    /**
     * 批量矫正文本框
     *
     * @param boxList
     * @param srcMat
     * @param manager
     * @return
     */
    private List<Image> batchAlign(List<OcrBox> boxList, Mat srcMat, NDManager manager) {
        List<Image> imageList = new ArrayList<>(boxList.size());
        for (int i = 0; i < boxList.size(); i++) {
            //透视变换 + 裁剪
            Image subImg = OcrUtils.transformAndCrop(srcMat, boxList.get(i));
            //ImageUtils.saveImage(subImg, i + "crop.png", "build/output");
            //高宽比 > 1.5
            if (subImg.getHeight() * 1.0 / subImg.getWidth() > 1.5) {
                //旋转图片90度
                subImg = OcrUtils.rotateImg(manager, subImg);
                //ImageUtils.saveImage(subImg, i + "rotate.png", "build/output");
            }
            imageList.add(subImg);
        }
        return imageList;
    }

    /**
     * 批量矫正文本框
     *
     * @param itemList
     * @param srcMat
     * @param manager
     * @return
     */
    private List<Image> batchAlignWithDirection(List<OcrItem> itemList, Mat srcMat, NDManager manager) {
        List<Image> imageList = new ArrayList<>(itemList.size());
        for (OcrItem ocrItem : itemList) {
            //放射变换+裁剪
            Image subImage = OcrUtils.transformAndCrop(srcMat, ocrItem.getOcrBox());
            //ImageUtils.saveImage(subImage, UUID.randomUUID().toString() + "_aaa.png", "build/output");
            //纠正文本框
            subImage = OcrUtils.rotateImg(subImage, ocrItem.getAngle());
            imageList.add(subImage);
        }
        return imageList;
    }


//    private RotatedBox recognize(OcrBox box,Mat srcMat,Predictor<Image, String> recPredictor,NDManager manager){
//        try {
//            //透视变换 + 裁剪
//            Image subImg = OcrUtils.transformAndCrop(srcMat, box);
//            //ImageUtils.saveImage(subImg, i + "crop.png", "build/output");
//            //高宽比 > 1.5
//            if (subImg.getHeight() * 1.0 / subImg.getWidth() > 1.5) {
//                //旋转图片90度
//                subImg = OcrUtils.rotateImg(manager, subImg);
//                //ImageUtils.saveImage(subImg, i + "rotate.png", "build/output");
//            }
//            String name = recPredictor.predict(subImg);
//            ((Mat)subImg.getWrappedImage()).release();
//            NDArray pointsArray = manager.create(box.toFloatArray());
//            return new RotatedBox(pointsArray, name);
//        } catch (Exception e) {
//            throw new OcrException("OCR检测错误", e);
//        }
//    }


    /**
     * 后处理：排序,分行
     *
     * @param rotatedBoxes
     */
    private OcrInfo postProcessOcrResult(List<RotatedBox> rotatedBoxes, OcrRecOptions ocrRecOptions) {
        //不分行
        if (!ocrRecOptions.isEnableLineSplit()) {
            return OcrUtils.convertRotatedBoxesToOcrItems(rotatedBoxes);
        }
        //Y坐标升序排序
        List<RotatedBox> initList = new ArrayList<>();
        for (RotatedBox result : rotatedBoxes) {
            initList.add(result);
        }
        Collections.sort(initList);
        //多行文本框的集合
        List<ArrayList<RotatedBoxCompX>> lines = new ArrayList<>();
        List<RotatedBoxCompX> line = new ArrayList<>();
        RotatedBoxCompX firstBox = new RotatedBoxCompX(initList.get(0).getBox(), initList.get(0).getText());
        line.add(firstBox);
        lines.add((ArrayList) line);
        //分行判断
        for (int i = 1; i < initList.size(); i++) {
            RotatedBoxCompX tmpBox = new RotatedBoxCompX(initList.get(i).getBox(), initList.get(i).getText());
            float y1 = firstBox.getBox().toFloatArray()[1];
            float y2 = tmpBox.getBox().toFloatArray()[1];
            float dis = Math.abs(y2 - y1);
            if (dis < 20) { // 认为是同 1 行  - Considered to be in the same line
                line.add(tmpBox);
            } else { // 换行 - Line break
                firstBox = tmpBox;
                Collections.sort(line);
                line = new ArrayList<>();
                line.add(firstBox);
                lines.add((ArrayList) line);
            }
        }
        return OcrUtils.convertToOcrInfo(lines);
    }


    @Override
    public void recognizeAndDraw(String imagePath, String outputPath, int fontSize, OcrRecOptions options) {
        if (!FileUtils.isFileExists(imagePath)) {
            throw new OcrException("图像文件不存在");
        }
        try {
            Image img = ImageFactory.getInstance().fromFile(Paths.get(imagePath));
            OcrInfo ocrInfo = recognize(img, options);
            if (Objects.isNull(ocrInfo) || Objects.isNull(ocrInfo.getLineList()) || ocrInfo.getLineList().isEmpty()) {
                throw new OcrException("未检测到文字");
            }
            Mat wrappedImage = (Mat) img.getWrappedImage();
            BufferedImage bufferedImage = OpenCVUtils.mat2Image(wrappedImage);
            OcrUtils.drawRectWithText(bufferedImage, ocrInfo, fontSize);
            ImageUtils.saveImage(bufferedImage, outputPath);
            wrappedImage.release();
        } catch (IOException e) {
            throw new OcrException(e);
        }
    }

    @Override
    public OcrInfo recognize(BufferedImage image, OcrRecOptions options) {
        if (!ImageUtils.isImageValid(image)) {
            throw new OcrException("图像无效");
        }
        Image img = ImageFactory.getInstance().fromImage(OpenCVUtils.image2Mat(image));
        OcrInfo ocrInfo = recognize(img, options);
        ((Mat) img.getWrappedImage()).release();
        return ocrInfo;
    }

    @Override
    public OcrInfo recognize(byte[] imageData, OcrRecOptions options) {
        if (Objects.isNull(imageData)) {
            throw new OcrException("图像无效");
        }
        try {
            BufferedImage image = ImageIO.read(new ByteArrayInputStream(imageData));
            return recognize(image, options);
        } catch (IOException e) {
            throw new OcrException("错误的图像", e);
        }
    }

    @Override
    public BufferedImage recognizeAndDraw(BufferedImage sourceImage, int fontSize, OcrRecOptions options) {
        if (!ImageUtils.isImageValid(sourceImage)) {
            throw new OcrException("图像无效");
        }
        Image img = ImageFactory.getInstance().fromImage(OpenCVUtils.image2Mat(sourceImage));
        OcrInfo ocrInfo = recognize(img, options);
        if (Objects.isNull(ocrInfo) || Objects.isNull(ocrInfo.getLineList()) || ocrInfo.getLineList().isEmpty()) {
            throw new OcrException("未检测到文字");
        }
        OcrUtils.drawRectWithText(sourceImage, ocrInfo, fontSize);
        return sourceImage;
    }

    @Override
    public String recognizeAndDrawToBase64(byte[] imageData, int fontSize, OcrRecOptions options) {
        if (Objects.isNull(imageData)) {
            throw new OcrException("图像无效");
        }
        OcrInfo ocrInfo = recognize(imageData, options);
        if (Objects.isNull(ocrInfo) || Objects.isNull(ocrInfo.getLineList()) || ocrInfo.getLineList().isEmpty()) {
            throw new OcrException("未检测到文字");
        }
        try {
            BufferedImage sourceImage = ImageIO.read(new ByteArrayInputStream(imageData));
            OcrUtils.drawRectWithText(sourceImage, ocrInfo, fontSize);
            return ImgUtil.toBase64(sourceImage, "png");
        } catch (IOException e) {
            throw new OcrException("导出图片失败", e);
        }
    }

    @Override
    public OcrInfo recognizeAndDraw(byte[] imageData, int fontSize, OcrRecOptions options) {
        if (Objects.isNull(imageData)) {
            throw new OcrException("图像无效");
        }
        OcrInfo ocrInfo = recognize(imageData, options);
        if (Objects.isNull(ocrInfo) || Objects.isNull(ocrInfo.getLineList()) || ocrInfo.getLineList().isEmpty()) {
            throw new OcrException("未检测到文字");
        }
        try {
            BufferedImage sourceImage = ImageIO.read(new ByteArrayInputStream(imageData));
            OcrUtils.drawRectWithText(sourceImage, ocrInfo, fontSize);
            ocrInfo.setBase64Img(ImgUtil.toBase64(sourceImage, "png"));
            return ocrInfo;
        } catch (IOException e) {
            throw new OcrException("导出图片失败", e);
        }
    }

    @Override
    public List<OcrInfo> batchRecognize(List<BufferedImage> imageList, OcrRecOptions options) {
        List<Image> djlImageList = new ArrayList<>(imageList.size());
        try {
            for (BufferedImage bufferedImage : imageList) {
                djlImageList.add(ImageFactory.getInstance().fromImage(OpenCVUtils.image2Mat(bufferedImage)));
            }
            return batchRecognizeDJLImage(djlImageList, options);
        } catch (Exception e) {
            throw new OcrException(e);
        } finally {
            djlImageList.forEach(image -> ((Mat) image.getWrappedImage()).release());
        }
    }

    @Override
    public List<OcrInfo> batchRecognizeDJLImage(List<Image> imageList, OcrRecOptions options) {
        if (Objects.isNull(textDetModel)) {
            throw new OcrException("textDetModel is null");
        }
        OcrRecOptions ocrRecOptions = options;
        if (Objects.isNull(options)) {
            ocrRecOptions = new OcrRecOptions();
        }
        if (CollectionUtils.isEmpty(imageList)) {
            throw new OcrException("imageList is empty");
        }
        //检测文本
        List<List<OcrBox>> boxeList = textDetModel.batchDetectDJLImage(imageList);
        if (CollectionUtils.isEmpty(boxeList) || boxeList.size() != imageList.size()) {
            throw new OcrException("未检测到文本");
        }
        Predictor<Image, String> predictor = null;
        List<OcrInfo> ocrInfoList = new ArrayList<OcrInfo>();
        try (NDManager manager = NDManager.newBaseManager()) {
            predictor = recPredictorPool.borrowObject();
            List<Image> allImageAlignList = new ArrayList<Image>();
            //检测方向
            if (ocrRecOptions.isEnableDirectionCorrect()) {
                if (Objects.isNull(directionModel)) {
                    throw new OcrException("请配置方向模型");
                }
                List<Mat> matList = imageList.stream()
                        .map(image -> (Mat) image.getWrappedImage())
                        .collect(Collectors.toList());
                List<List<OcrItem>> ocrItemList = directionModel.batchDetect(boxeList, matList);
                if (CollectionUtils.isEmpty(ocrItemList) || ocrItemList.size() != imageList.size()) {
                    throw new OcrException("方向检测失败");
                }
                allImageAlignList = new ArrayList<Image>();
                for (int i = 0; i < ocrItemList.size(); i++) {
                    Mat srcMat = (Mat) imageList.get(i).getWrappedImage();
                    List<Image> imageAlignList = batchAlignWithDirection(ocrItemList.get(i), srcMat, manager);
//                    for(int j = 0; j < imageAlignList.size(); j++){
//                        ImageUtils.saveImage(imageAlignList.get(j),"dir-"+i+"-"+j+".png","/Users/xxx/Downloads/testing33");
//                    }
                    allImageAlignList.addAll(imageAlignList);
                }
            } else {
                for (int i = 0; i < boxeList.size(); i++) {
                    Mat srcMat = (Mat) imageList.get(i).getWrappedImage();
                    List<Image> imageAlignList = batchAlign(boxeList.get(i), srcMat, manager);
//                    for(int j = 0; j < imageAlignList.size(); j++){
//                        ImageUtils.saveImage(imageAlignList.get(j),i+"-"+j+".png","/Users/xxx/Downloads/testing33");
//                    }
                    allImageAlignList.addAll(imageAlignList);
                }
            }
            List<String> textList = batchRecognize(allImageAlignList);
            int textIndex = 0;
            for (int i = 0; i < boxeList.size(); i++) {
                List<RotatedBox> rotatedBoxes = new ArrayList<>();
                for (int j = 0; j < boxeList.get(i).size(); j++) {
                    if (textIndex >= textList.size()) {
                        throw new OcrException("识别失败: 第" + i + "张图片, 第" + j + "个文本块，未识别到文本");
                    }
                    OcrBox box = boxeList.get(i).get(j);
                    NDArray pointsArray = manager.create(box.toFloatArray());
                    rotatedBoxes.add(new RotatedBox(pointsArray, textList.get(textIndex)));
                    textIndex++;
                }
                OcrInfo ocrInfo = postProcessOcrResult(rotatedBoxes, ocrRecOptions);
                ocrInfoList.add(ocrInfo);
            }
            return ocrInfoList;
        } catch (Exception e) {
            throw new OcrException("OCR检测错误", e);
        } finally {
            if (predictor != null) {
                try {
                    recPredictorPool.returnObject(predictor); //归还
                } catch (Exception e) {
                    log.warn("归还Predictor失败", e);
                    try {
                        predictor.close(); // 归还失败才销毁
                    } catch (Exception ex) {
                        log.error("关闭Predictor失败", ex);
                    }
                }
            }
        }
    }

    private List<String> batchRecognize(List<Image> imageAlignList) {
        Predictor<Image, String> predictor = null;
        try {
            predictor = recPredictorPool.borrowObject();
            List<String> textList = predictor.batchPredict(imageAlignList);
            imageAlignList.forEach(subImg -> ((Mat) subImg.getWrappedImage()).release());
            return textList;
        } catch (Exception e) {
            throw new OcrException("OCR检测错误", e);
        } finally {
            if (predictor != null) {
                try {
                    recPredictorPool.returnObject(predictor); //归还
                } catch (Exception e) {
                    log.warn("归还Predictor失败", e);
                    try {
                        predictor.close(); // 归还失败才销毁
                    } catch (Exception ex) {
                        log.error("关闭Predictor失败", ex);
                    }
                }
            }
        }
    }

    @Override
    public void setTextDetModel(OcrCommonDetModel detModel) {
        this.textDetModel = detModel;
    }

    @Override
    public OcrCommonDetModel getTextDetModel() {
        return textDetModel;
    }

    @Override
    public void setDirectionModel(OcrDirectionModel directionModel) {
        this.directionModel = directionModel;
    }

    @Override
    public OcrDirectionModel getDirectionModel() {
        return directionModel;
    }


    public GenericObjectPool<Predictor<Image, String>> getRecPredictorPool() {
        return recPredictorPool;
    }

    @Override
    public void close() throws Exception {
        try {
            if (recPredictorPool != null) {
                recPredictorPool.close();
            }
        } catch (Exception e) {
            log.warn("关闭 predictorPool 失败", e);
        }
        try {
            if (recognitionModel != null) {
                recognitionModel.close();
            }
        } catch (Exception e) {
            log.warn("关闭 model 失败", e);
        }
    }
}