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

import cn.hutool.core.bean.BeanUtil;
import cn.hutool.core.util.StrUtil;
import com.alibaba.fastjson2.JSONObject;
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.pavis.admin.aigc.enums.EmbedStoreEnum;
import com.pavis.admin.aigc.mapper.EmbedStoreMapper;
import com.pavis.admin.aigc.model.entity.EmbedStoreDO;
import com.pavis.admin.aigc.model.resp.EmbedStoreResp;
import com.pavis.admin.aigc.service.EmbedStoreService;
import dev.langchain4j.community.store.embedding.redis.RedisEmbeddingStore;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.milvus.MilvusEmbeddingStore;
import dev.langchain4j.store.embedding.pgvector.PgVectorEmbeddingStore;
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.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author tycoding
 * @since 2024/10/28
 */
@Slf4j
@Component
public class EmbeddingStoreFactory {

    @Autowired
    private EmbedStoreService aigcEmbedStoreService;

    @Autowired
    private EmbedStoreMapper aigcEmbedStoreMapper;

    private final List<EmbedStoreResp> modelStore = new ArrayList<>();
    private final Map<String, EmbeddingStore<TextSegment>> embedStoreMap = new ConcurrentHashMap<>();

    @Async
    @PostConstruct
    public void init() {
        modelStore.clear();
        List<EmbedStoreDO> embedStoreDOS = aigcEmbedStoreMapper.selectList(new QueryWrapper<>());
        List<EmbedStoreResp> list = BeanUtil.copyToList(embedStoreDOS, EmbedStoreResp.class);
        // List<EmbedStoreResp> list = aigcEmbedStoreService.list(new EmbedStoreQuery(),new SortQuery());
        list.forEach(embed -> {
            JSONObject endpoints = JSONObject.parseObject(embed.getEndpoints());
            embed.setHost(endpoints.get("host").toString());
            embed.setPort((Integer)endpoints.get("port"));
            JSONObject databaseConfig = JSONObject.parseObject(embed.getDatabaseConfig());
            embed.setDatabaseName(databaseConfig.get("database").toString());
            embed.setTableName(databaseConfig.get("table").toString());
            JSONObject authConfig = JSONObject.parseObject(embed.getAuthConfig());
            embed.setUsername(authConfig.get("username").toString());
            embed.setPassword(authConfig.get("password").toString());
            JSONObject vectorConfig = JSONObject.parseObject(embed.getVectorConfig());
            embed.setDimension((Integer)vectorConfig.get("dimension"));
            try {
                if (EmbedStoreEnum.REDIS.name().equalsIgnoreCase(embed.getProvider())) {
                    RedisEmbeddingStore.Builder builder = RedisEmbeddingStore.builder()
                        .host(embed.getHost())
                        .port(embed.getPort())
                        .indexName(embed.getDatabaseName())
                        .dimension(embed.getDimension());
                    if (StrUtil.isNotBlank(embed.getUsername()) && StrUtil.isNotBlank(embed.getPassword())) {
                        builder.user(embed.getUsername()).password(embed.getPassword());
                    }
                    EmbeddingStore<TextSegment> store = builder.build();
                    embedStoreMap.put(embed.getId().toString(), store);
                }
                if (EmbedStoreEnum.PGVECTOR.name().equalsIgnoreCase(embed.getProvider())) {
                    EmbeddingStore<TextSegment> store = PgVectorEmbeddingStore.builder()
                        .host(embed.getHost())
                        .port(embed.getPort())
                        .database(embed.getDatabaseName())
                        .dimension(embed.getDimension())
                        .user(embed.getUsername())
                        .password(embed.getPassword())
                        .table(embed.getTableName())
                        .indexListSize(1)
                        .useIndex(true)
                        .createTable(true)
                        .dropTableFirst(false)
                        .build();
                    embedStoreMap.put(embed.getId().toString(), store);
                }
                if (EmbedStoreEnum.MILVUS.name().equalsIgnoreCase(embed.getProvider())) {
                    EmbeddingStore<TextSegment> store = MilvusEmbeddingStore.builder()
                        .host(embed.getHost())
                        .port(embed.getPort())
                        .databaseName(embed.getDatabaseName())
                        .dimension(embed.getDimension())
                        .username(embed.getUsername())
                        .password(embed.getPassword())
                        .collectionName(embed.getTableName())
                        .build();
                    embedStoreMap.put(embed.getId().toString(), store);
                }
                modelStore.add(embed);
            } catch (Exception e) {
                e.printStackTrace();
                log.error("向量数据库初始化失败：[{}] --- [{}]，数据库配置信息：[{}]", embed.getName(), embed.getProvider(), embed);
            }
        });

        modelStore.forEach(i -> log.info("已成功注册Embedding Store：{}， 配置信息：{}", i.getProvider(), i));
    }

    public EmbeddingStore<TextSegment> getEmbeddingStore(String embeddingId) {
        return embedStoreMap.get(embeddingId);
    }

    public boolean containsEmbeddingStore(String embeddingId) {
        return embedStoreMap.containsKey(embeddingId);
    }
}
