package tech.aiflowy.common.ai;

import tech.aiflowy.common.ai.util.LLMUtil;
import tech.aiflowy.common.options.SysOptions;
import tech.aiflowy.common.util.StringUtil;
import com.agentsflex.core.llm.ChatContext;
import com.agentsflex.core.llm.Llm;
import com.agentsflex.core.llm.StreamResponseListener;
import com.agentsflex.core.llm.response.AiMessageResponse;
import com.agentsflex.core.prompt.TextPrompt;
import com.agentsflex.core.store.DocumentStore;
import com.agentsflex.store.aliyun.AliyunVectorStore;
import com.agentsflex.store.aliyun.AliyunVectorStoreConfig;
import com.agentsflex.store.milvus.MilvusVectorStore;
import com.agentsflex.store.milvus.MilvusVectorStoreConfig;
import com.agentsflex.store.qcloud.QCloudVectorStore;
import com.agentsflex.store.qcloud.QCloudVectorStoreConfig;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

public class EmbeddingManager {

    private static final Logger logger = LoggerFactory.getLogger(EmbeddingManager.class);
    private static final EmbeddingManager manager = new EmbeddingManager();

    public static EmbeddingManager getInstance() {
        return manager;
    }

    private ExecutorService sseExecutor = Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors());

    public Llm getEmbeddingModel() {
        String modelOfEmbedding = SysOptions.get("model_of_embedding");
        return LLMUtil.getLlmByType(modelOfEmbedding);
    }

    public DocumentStore getVectorStore() {
        String vectorStoreType = SysOptions.get("vectorstore_type");
        if (StringUtil.noText(vectorStoreType)) {
            return null;
        }

        DocumentStore store = null;

        if ("aliyun".equalsIgnoreCase(vectorStoreType)) {
            AliyunVectorStoreConfig storeConfig = new AliyunVectorStoreConfig();
            storeConfig.setDefaultCollectionName(SysOptions.get("vectorstore_default_collection"));
            storeConfig.setApiKey(SysOptions.get("aliyun_vdb_api_key"));
            storeConfig.setEndpoint(SysOptions.get("aliyun_vdb_endpoint"));
            storeConfig.setDatabase(SysOptions.get("aliyun_vdb_database"));
            store = new AliyunVectorStore(storeConfig);
        }
        //腾讯云
        else if ("qcloud".equalsIgnoreCase(vectorStoreType)) {
            QCloudVectorStoreConfig storeConfig = new QCloudVectorStoreConfig();
            storeConfig.setDefaultCollectionName(SysOptions.get("vectorstore_default_collection"));
            storeConfig.setHost(SysOptions.get("qcloud_vdb_endpoint"));
            storeConfig.setApiKey(SysOptions.get("qcloud_vdb_api_key"));
            storeConfig.setAccount(SysOptions.get("qcloud_vdb_username"));
            storeConfig.setDatabase(SysOptions.get("qcloud_vdb_database"));
            store = new QCloudVectorStore(storeConfig);
        }
        //Milvus 数据库
        else if ("milvus".equalsIgnoreCase(vectorStoreType)) {
            MilvusVectorStoreConfig storeConfig = new MilvusVectorStoreConfig();
            storeConfig.setDatabaseName(SysOptions.get("milvus_database"));
            storeConfig.setDefaultCollectionName(SysOptions.get("vectorstore_default_collection"));
            storeConfig.setUri(SysOptions.get("milvus_endpoint"));
            storeConfig.setToken(SysOptions.get("milvus_apikey"));
            store = new MilvusVectorStore(storeConfig);
        }

        Llm embeddingModel = getEmbeddingModel();
        if (store != null && embeddingModel != null) {
            store.setEmbeddingModel(embeddingModel);
        }

        return store;
    }


    public SseEmitter sseEmitter(String prompt) {
        return sseEmitter(new TextPrompt(prompt));
    }

    public SseEmitter sseEmitter(TextPrompt prompt) {
        SseEmitter emitter = new SseEmitter((long) (1000 * 60 * 2));
        sseExecutor.execute(() -> {
            Llm llm = getEmbeddingModel();
            if (llm == null) {
                try {
                    emitter.send(SseEmitter.event().data("AI 大模型未配置正确"));
                    emitter.complete();
                } catch (IOException e) {
                    logger.error(e.toString(), e);
                    emitter.completeWithError(e);
                }
                return;
            }
            llm.chatStream(prompt, new StreamResponseListener() {
                @Override
                public void onMessage(ChatContext chatContext, AiMessageResponse aiMessageResponse) {
                    try {
                        String content = aiMessageResponse.getMessage().getContent();
                        System.out.println("ai send >>>>>>" + content + "<<<<<<");
                        emitter.send(SseEmitter.event().data(content));
                    } catch (IOException e) {
                        logger.error(e.toString(), e);
                        emitter.completeWithError(e);
                    }
                }

                @Override
                public void onStop(ChatContext context) {
                    emitter.complete();
                }
            });
        });
        return emitter;
    }

    public ExecutorService getSseExecutor() {
        return sseExecutor;
    }

    public void setSseExecutor(ExecutorService sseExecutor) {
        this.sseExecutor = sseExecutor;
    }
}