package com.pavis.admin.aigc.core.llm.provider;

import com.pavis.admin.aigc.model.resp.EmbedStoreResp;
import com.pavis.admin.aigc.model.resp.KnowledgeResp;
import com.pavis.admin.aigc.model.resp.ModelResp;
import com.pavis.admin.aigc.service.EmbedStoreService;
import com.pavis.admin.aigc.service.KnowledgeService;
import com.pavis.admin.aigc.service.ModelService;
import jakarta.annotation.PostConstruct;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.scheduling.annotation.Async;
import org.springframework.stereotype.Component;

import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;

/**
 * @author tycoding
 * @since 2024/10/29
 */
@Slf4j
@Component
public class KnowledgeStoreFactory {

    @Autowired
    private KnowledgeService knowledgeService;
    @Autowired
    private ModelService modelService;
    @Autowired
    private EmbedStoreService embedStoreService;

    private final Map<String, KnowledgeResp> knowledgeMap = new ConcurrentHashMap<>();

    @Async
    @PostConstruct
    public void init() {
        knowledgeMap.clear();
        List<KnowledgeResp> list = knowledgeService.selectKnowledge();
        Map<Long, List<ModelResp>> modelMap = modelService.selModelAndSecrt()
            .stream()
            .collect(Collectors.groupingBy(ModelResp::getId));
        Map<Long, List<EmbedStoreResp>> storeMap = embedStoreService.selEmbedStore()
            .stream()
            .collect(Collectors.groupingBy(EmbedStoreResp::getId));
        list.forEach(know -> {
            if (know.getEmbedModelId() != null) {
                List<ModelResp> models = modelMap.get(know.getEmbedModelId());
                know.setEmbedModel(models == null ? null : models.get(0));
            }
            if (know.getEmbedStoreId() != null) {
                List<EmbedStoreResp> stores = storeMap.get(know.getEmbedStoreId());
                know.setEmbedStore(stores == null ? null : stores.get(0));
            }
            knowledgeMap.put(know.getId().toString(), know);
        });
    }

    public KnowledgeResp getKnowledge(Long knowledgeId) {
        return knowledgeMap.get(knowledgeId.toString());
    }

    public boolean containsKnowledge(Long knowledgeId) {
        return knowledgeMap.containsKey(knowledgeId.toString());
    }
}
