
package tech.easyflow.ai.entity;

import cn.hutool.core.util.StrUtil;
import com.easyagents.core.model.chat.ChatModel;
import com.easyagents.core.model.embedding.EmbeddingModel;
import com.easyagents.core.model.rerank.RerankModel;
import com.easyagents.core.store.VectorData;
import com.easyagents.embedding.ollama.OllamaEmbeddingConfig;
import com.easyagents.embedding.ollama.OllamaEmbeddingModel;
import com.easyagents.embedding.openai.OpenAIEmbeddingConfig;
import com.easyagents.embedding.openai.OpenAIEmbeddingModel;
import com.easyagents.llm.deepseek.DeepseekChatModel;
import com.easyagents.llm.deepseek.DeepseekConfig;
import com.easyagents.llm.ollama.OllamaChatConfig;
import com.easyagents.llm.ollama.OllamaChatModel;
import com.easyagents.llm.openai.OpenAIChatConfig;
import com.easyagents.llm.openai.OpenAIChatModel;
import com.easyagents.rerank.DefaultRerankModel;
import com.easyagents.rerank.DefaultRerankModelConfig;
import com.easyagents.rerank.gitee.GiteeRerankModel;
import com.easyagents.rerank.gitee.GiteeRerankModelConfig;
import com.mybatisflex.annotation.RelationManyToOne;
import com.mybatisflex.annotation.Table;
import tech.easyflow.ai.entity.base.ModelBase;
import tech.easyflow.common.util.StringUtil;
import tech.easyflow.common.web.exceptions.BusinessException;

/**
 * 实体类。
 *
 * @author michael
 * @since 2024-08-23
 */

@Table("tb_model")
public class Model extends ModelBase {

    @RelationManyToOne(selfField = "providerId", targetField = "id")
    private ModelProvider modelProvider;

    /**
     * 模型类型
     */
    public final static String[] MODEL_TYPES = {"chatModel", "embeddingModel", "rerankModel"};


    public ModelProvider getModelProvider() {
        return modelProvider;
    }

    public void setModelProvider(ModelProvider modelProvider) {
        this.modelProvider = modelProvider;
    }

    public ChatModel toChatModel() {
        String providerType = modelProvider.getProviderType();
        if (StringUtil.noText(providerType)) {
            return null;
        }
        switch (providerType.toLowerCase()) {
            case "ollama":
                OllamaChatConfig ollamaChatConfig = new OllamaChatConfig();
                ollamaChatConfig.setEndpoint(checkAndGetEndpoint());
                if (StringUtil.hasText(getApiKey())) {
                    ollamaChatConfig.setApiKey(checkAndGetApiKey());
                }
                ollamaChatConfig.setModel(checkAndGetModelName());
                ollamaChatConfig.setProvider(getModelProvider().getProviderName());
                ollamaChatConfig.setSupportImageBase64Only(getSupportImageB64Only());
                return new OllamaChatModel(ollamaChatConfig);
            case "deepseek":
                DeepseekConfig deepseekConfig = new DeepseekConfig();
                deepseekConfig.setProvider(getModelProvider().getProviderName());
                deepseekConfig.setEndpoint(checkAndGetEndpoint());
                deepseekConfig.setApiKey(checkAndGetApiKey());
                deepseekConfig.setModel(checkAndGetModelName());
                deepseekConfig.setRequestPath(checkAndGetRequestPath());
                deepseekConfig.setSupportThinking(Boolean.TRUE);
                deepseekConfig.setThinkingProtocol("deepseek");
                deepseekConfig.setNeedReasoningContentForToolMessage(Boolean.TRUE);
                deepseekConfig.setSupportImageBase64Only(getSupportImageB64Only());
                if (getSupportToolMessage() != null) {
                    deepseekConfig.setSupportToolMessage(getSupportToolMessage());
                }
                return new DeepseekChatModel(deepseekConfig);
            default:
                OpenAIChatConfig openAIChatConfig = new OpenAIChatConfig();
                openAIChatConfig.setProvider(getModelProvider().getProviderName());
                openAIChatConfig.setEndpoint(checkAndGetEndpoint());
                openAIChatConfig.setApiKey(checkAndGetApiKey());
                openAIChatConfig.setModel(checkAndGetModelName());
                openAIChatConfig.setRequestPath(checkAndGetRequestPath());
                openAIChatConfig.setSupportImageBase64Only(getSupportImageB64Only());
                if (getSupportToolMessage() != null) {
                    openAIChatConfig.setSupportToolMessage(getSupportToolMessage());
                }
                return new OpenAIChatModel(openAIChatConfig);
        }
    }

    public RerankModel toRerankModel() {
        switch (modelProvider.getProviderType().toLowerCase()) {
            case "gitee":
                GiteeRerankModelConfig giteeRerankModelConfig = new GiteeRerankModelConfig();
                giteeRerankModelConfig.setProvider(getModelProvider().getProviderName());
                giteeRerankModelConfig.setApiKey(checkAndGetApiKey());
                giteeRerankModelConfig.setEndpoint(checkAndGetEndpoint());
                giteeRerankModelConfig.setModel(checkAndGetModelName());
                giteeRerankModelConfig.setRequestPath(checkAndGetRequestPath());
                return new GiteeRerankModel(giteeRerankModelConfig);
            default:
                DefaultRerankModelConfig defaultRerankModelConfig = new DefaultRerankModelConfig();
                defaultRerankModelConfig.setProvider(getModelProvider().getProviderName());
                defaultRerankModelConfig.setApiKey(checkAndGetApiKey());
                defaultRerankModelConfig.setEndpoint(checkAndGetEndpoint());
                defaultRerankModelConfig.setRequestPath(checkAndGetRequestPath());
                defaultRerankModelConfig.setModel(checkAndGetModelName());
                return new DefaultRerankModel(defaultRerankModelConfig);
        }
    }

    public EmbeddingModel toEmbeddingModel() {
        String providerType = modelProvider.getProviderType();
        if (StringUtil.noText(providerType)) {
            return null;
        }
        try {
            switch (providerType.toLowerCase()) {
                case "ollama":
                    OllamaEmbeddingConfig ollamaEmbeddingConfig = new OllamaEmbeddingConfig();
                    ollamaEmbeddingConfig.setProvider(getModelProvider().getProviderName());
                    ollamaEmbeddingConfig.setEndpoint(checkAndGetEndpoint());
                    ollamaEmbeddingConfig.setApiKey(getApiKey());
                    ollamaEmbeddingConfig.setModel(checkAndGetModelName());
                    ollamaEmbeddingConfig.setRequestPath(getRequestPath());
                    return new OllamaEmbeddingModel(ollamaEmbeddingConfig);
                default:
                    OpenAIEmbeddingConfig openAIEmbeddingConfig = new OpenAIEmbeddingConfig();
                    openAIEmbeddingConfig.setProvider(getModelProvider().getProviderName());
                    openAIEmbeddingConfig.setEndpoint(checkAndGetEndpoint());
                    openAIEmbeddingConfig.setApiKey(checkAndGetApiKey());
                    openAIEmbeddingConfig.setModel(checkAndGetModelName());
                    openAIEmbeddingConfig.setRequestPath(checkAndGetRequestPath());
                    return new OpenAIEmbeddingModel(openAIEmbeddingConfig);
            }
        } catch (Exception e) {
            throw new BusinessException("向量模型配置失败：" + e.getMessage());
        }

    }

    /**
     * 获取模型向量的维度
     *
     * @return
     */
    public static int getEmbeddingDimension(EmbeddingModel embeddingModel) {
        if (embeddingModel == null) {
            throw new BusinessException("embeddingModel不能为空");
        }
        VectorData vectorData = embeddingModel.embed("测试向量维度");
        return vectorData.getVector().length;
    }

    public String checkAndGetRequestPath() {
        if (StrUtil.isEmpty(getRequestPath())) {
            throw new BusinessException("请求地址不能为空");
        }
        return getRequestPath();
    }

    public String checkAndGetApiKey() {
        if (StrUtil.isEmpty(getApiKey())) {
            throw new BusinessException("API 密钥不能为空");
        }
        return getApiKey();
    }

    public String checkAndGetEndpoint() {
        if (StrUtil.isEmpty(getEndpoint())) {
            throw new BusinessException("API 地址不能为空");
        }
        return getEndpoint();
    }

    public String checkAndGetModelName() {
        if (StrUtil.isEmpty(getModelName())) {
            throw new BusinessException("模型名称不能为空");
        }
        return getModelName();
    }

}
