fu-hsi-service/src/main/java/com/supervision/police/service/impl/FileOcrProcessServiceImpl.java

package com.supervision.police.service.impl;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONUtil;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import com.supervision.minio.domain.MinioFile;
import com.supervision.minio.service.MinioService;
import com.supervision.police.domain.FileOcrProcess;
import com.supervision.police.domain.NoteRecord;
import com.supervision.police.dto.OCRReqDTO;
import com.supervision.police.dto.OCRResDTO;
import com.supervision.police.dto.RecordFileDTO;
import com.supervision.police.service.FileOcrProcessService;
import com.supervision.police.mapper.FileOcrProcessMapper;
import com.supervision.police.service.NoteRecordService;
import com.supervision.police.service.OCRService;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.scheduling.annotation.Async;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Objects;
import java.util.stream.Collectors;

/**
* @author Administrator
* @description 针对表【file_ocr_process(文件ocr识别进度表)】的数据库操作Service实现
* @createDate 2024-08-30 17:35:23
*/
@Slf4j
@Service
@RequiredArgsConstructor
public class FileOcrProcessServiceImpl extends ServiceImpl<FileOcrProcessMapper, FileOcrProcess>
    implements FileOcrProcessService{

    @Value("${ocr.pool.max-size:20}")
    private Integer poolMaxSize;

    private final OCRService ocrService;

    private final MinioService minioService;

    private final NoteRecordService noteRecordService;

    @Override
    @Transactional(transactionManager = "dataSourceTransactionManager",rollbackFor = Exception.class)
    public List<FileOcrProcess> submitOCR(List<String> fileIdList) {

        if (CollUtil.isEmpty(fileIdList)){
            log.info("submitOCR:fileIds为空。提交ocr任务...");
            return new ArrayList<>(1);
        }
        List<FileOcrProcess> processList = new ArrayList<>();
        for (String fileId : fileIdList) {
            FileOcrProcess fileOcrProcess = new FileOcrProcess(fileId, -1);
            super.save(fileOcrProcess);
            processList.add(fileOcrProcess);
        }

        doOCRTask(processList);
        return processList;
    }

    @Async
    @Override
    public synchronized void doOCRTask(List<FileOcrProcess> fileOcrProcesses) {


        if (CollUtil.isEmpty(fileOcrProcesses)){
            log.info("asyncOcr:当前暂无识别的任务，结束...");
            return;
        }

        List<List<FileOcrProcess>> ocrTaskList = CollUtil.split(fileOcrProcesses, poolMaxSize);
        for (List<FileOcrProcess> ocrProcesses : ocrTaskList) {
            List<String> fileIdList = ocrProcesses.stream().map(FileOcrProcess::getFileId).collect(Collectors.toList());
            List<OCRReqDTO> ocrReqDTOS = buildOCRReqDTO(fileIdList);
            for (OCRReqDTO ocrReqDTO : ocrReqDTOS) {
                log.info("ocr:开始识别文件：{}", JSONUtil.toJsonStr(ocrReqDTO));
                this.updateOCrStatus(ocrReqDTO.getFile_ids(),0);

                try {
                    List<OCRResDTO> ocrRes = ocrService.ocr(ocrReqDTO);
                    log.info("ocr:识别结果：{}", JSONUtil.toJsonStr(ocrRes));
                    if (CollUtil.isNotEmpty(ocrRes)){
                        for (OCRResDTO ocrRe : ocrRes) {
                            if (Integer.valueOf(2).equals(ocrRe.getStatus())){
                                log.info("ocr:文件{}识别失败，原因：{}", ocrRe.getFile_id(), ocrRe.getError_msg());
                            }
                        }
                    }
                    ocrRes.forEach(this::updateByOcrRes);
                } catch (Exception e) {
                    log.error("远程调用ocr识别接口失败",e);
                    this.updateOCrStatus(ocrReqDTO.getFile_ids(),2);
                }

            }
        }
    }

    @Override
    public void doAllOCRTask() {
        List<FileOcrProcess> allFileOcrProcesses = pageListByStatus(-1, 99999);
        doOCRTask(allFileOcrProcesses);
    }

    @Override
    public List<FileOcrProcess> pageListByStatus(Integer status, Integer size) {
        return super.lambdaQuery().eq(FileOcrProcess::getStatus, status).page(new Page<>(1, size)).getRecords();
    }

    @Override
    public Integer countByStatus(Integer status) {
        return Math.toIntExact(super.lambdaQuery().eq(FileOcrProcess::getStatus, status).count());
    }

    @Override
    public Boolean updateOCrStatus(List<String> ocrIdList, Integer ocrStatus) {
        return super.lambdaUpdate().in(FileOcrProcess::getFileId, ocrIdList)
                .set(FileOcrProcess::getStatus, ocrStatus).update();
    }

    @Override
    public Boolean updateByOcrRes(OCRResDTO ocrResDTO) {
        return super.lambdaUpdate().eq(FileOcrProcess::getFileId, ocrResDTO.getFile_id())
                .set(FileOcrProcess::getOcrText, ocrResDTO.getOcr_text())
                .set(FileOcrProcess::getStatus, ocrResDTO.getStatus())
                .set(FileOcrProcess::getDrawImgId, ocrResDTO.getDraw_img_id())
                .update();
    }

    @Override
    public List<RecordFileDTO> queryFileList(List<String> fileIdList) {
        if (CollUtil.isEmpty(fileIdList)){
            return new ArrayList<>(1);
        }
        return super.baseMapper.queryFileList(fileIdList);
    }

    @Override
    public List<RecordFileDTO> queryListByRecordId(String recordId) {
        if (StrUtil.isEmpty(recordId)){
            return new ArrayList<>(1);
        }
        NoteRecord noteRecord = noteRecordService.getById(recordId);
        if (Objects.isNull(noteRecord) || StrUtil.isEmpty(noteRecord.getFileIds())){
            return new ArrayList<>(1);
        }
        return queryFileList(Arrays.stream(noteRecord.getFileIds().split(",")).toList());
    }


    private List<OCRReqDTO> buildOCRReqDTO(List<String> fileIdList){
        List<MinioFile> minioFiles = minioService.listMinioFile(fileIdList);
        return minioFiles.stream().collect(Collectors.groupingBy(MinioFile::getFileType))
                .entrySet().stream().map(entry ->
                        new OCRReqDTO(entry.getValue().stream().map(MinioFile::getId).collect(Collectors.toList()),entry.getKey()))
                .collect(Collectors.toList());

    }

}