package com.supervision.pdfqaserver.service.impl;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.StrUtil;
import com.supervision.pdfqaserver.cache.PromptCache;
import com.supervision.pdfqaserver.dto.*;
import com.supervision.pdfqaserver.service.TripleConversionPipeline;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.List;
@Slf4j
@Service
@RequiredArgsConstructor
public class TripleConversionPipelineImpl implements TripleConversionPipeline {

    private final OllamaChatModel ollamaChatModel;

    @Override
    public List<TruncateDTO> sliceDocuments(List<DocumentDTO> documents) {
        // 对pdfAnalysisOutputs进行排序
        List<DocumentDTO> documentDTOList = documents.stream().sorted(
                // 先对pageNo进行排序再对layoutOrder进行排序
                (o1, o2) -> {
                    if (o1.getPageNo().equals(o2.getPageNo())) {
                        return Integer.compare(o1.getLayoutOrder(), o2.getLayoutOrder());
                    }
                    return Integer.compare(o1.getPageNo(), o2.getPageNo());
                }
        ).toList();
        return null;
    }

    @Override
    public EREDTO doEre(TruncateDTO truncateDTO) {

        if (StrUtil.equals(truncateDTO.getLayoutType(),"0")){

            EREDTO eredto = doTextEre(truncateDTO);
            return eredto;
        }

        if (StrUtil.equals(truncateDTO.getLayoutType(),"1")){
            EREDTO eredto = doTableEre(truncateDTO);
            return eredto;
        }
        log.info("doEre:错误的布局类型: {}", truncateDTO.getLayoutType());
        return null;
    }

    private EREDTO doTextEre(TruncateDTO truncateDTO) {
        String prompt = PromptCache.promptMap.get(PromptCache.DOERE_TEXT);
        String formatted = String.format(prompt, truncateDTO.getContent());
        String response = ollamaChatModel.call(formatted);
        // todo:暂时不去处理异常返回

        return EREDTO.fromTextJson(response, truncateDTO.getId());
    }

    private EREDTO doTableEre(TruncateDTO truncateDTO) {
        String prompt = PromptCache.promptMap.get(PromptCache.DOERE_TABLE);
        String formatted = String.format(prompt, truncateDTO.getContent());
        String response = ollamaChatModel.call(formatted);
        // todo:暂时不去处理异常返回

        return EREDTO.fromTableJson(response, truncateDTO.getId());
    }

    /**
     * 合并实体关系抽取结果 主要是对实体和关系中的属性进行合并
     * @param eredtoList 实体关系抽取结果列表
     * @return
     */
    @Override
    public List<EREDTO> mergeEreResults(List<EREDTO> eredtoList) {
        List<EREDTO> merged = new ArrayList<>();
        if (CollUtil.isEmpty(eredtoList)){
            return merged;
        }
        for (EREDTO eredto : eredtoList) {
            List<EntityExtractionDTO> entities = eredto.getEntities();
            if (CollUtil.isNotEmpty(entities)){
                for (EntityExtractionDTO entity : entities) {
                    String e = entity.getEntity();
                    String name = entity.getName();
                    // entity.getEntity() 和 entity.getName() 完全相等看作是同一个数据
                }
            }
            List<RelationExtractionDTO> relations = eredto.getRelations();
            if (CollUtil.isNotEmpty(relations)){
                for (RelationExtractionDTO relation : relations) {
                    String source = relation.getSource();
                    String target = relation.getTarget();
                    String re = relation.getRelation();
                    // source和target,re完全相等看作是同一个数据
                }
            }
        }

        return null;
    }
}