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

import cn.hutool.core.util.StrUtil;
import com.pavis.admin.aigc.core.llm.exception.ServiceException;
import com.pavis.admin.aigc.enums.ChatErrorEnum;
import com.pavis.admin.aigc.enums.ProviderEnum;
import com.pavis.admin.aigc.model.resp.ModelResp;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.image.ImageModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.openai.OpenAiEmbeddingModel;
import dev.langchain4j.model.openai.OpenAiImageModel;
import dev.langchain4j.model.openai.OpenAiStreamingChatModel;
import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import java.time.Duration;

/**
 * @author tycoding
 * @since 2024-08-19 10:08
 */
@Slf4j
@Component
@AllArgsConstructor
public class OpenAIModelBuildHandler implements ModelBuildHandler {

    // private final LangChatProps props;

    /**
     * 合并处理支持OpenAI接口的模型
     */
    @Override
    public boolean whetherCurrentModel(ModelResp model) {
        String provider = model.getProvider();
        return ProviderEnum.OPENAI.name().equals(provider) || ProviderEnum.GEMINI.name()
            .equals(provider) || ProviderEnum.CLAUDE.name().equals(provider) || ProviderEnum.AZURE_OPENAI.name()
                .equals(provider) || ProviderEnum.DOUYIN.name().equals(provider) || ProviderEnum.YI.name()
                    .equals(provider) || ProviderEnum.SILICON.name().equals(provider) || ProviderEnum.DEEPSEEK.name()
                        .equals(provider) || ProviderEnum.SPARK.name().equals(provider);
    }

    @Override
    public boolean basicCheck(ModelResp model) {
        String apiKey = model.getApiKey();
        if (StrUtil.isBlank(apiKey)) {
            throw new ServiceException(ChatErrorEnum.API_KEY_IS_NULL.getErrorCode(), ChatErrorEnum.API_KEY_IS_NULL
                .getErrorDesc(model.getProvider().toUpperCase(), model.getType()));
        }
        return true;
    }

    @Override
    public StreamingChatLanguageModel buildStreamingChat(ModelResp model) {
        try {
            if (!whetherCurrentModel(model)) {
                return null;
            }
            if (!basicCheck(model)) {
                return null;
            }
            return OpenAiStreamingChatModel.builder()
                .apiKey(model.getApiKey())
                .baseUrl(model.getBaseUrl())
                .modelName(model.getModel())
                .maxTokens(model.getResponseLimit())
                .temperature(model.getTemperature())
                .logRequests(true)
                .logResponses(true)
                .topP(model.getTopP())
                .timeout(Duration.ofMinutes(10))
                .build();
        } catch (ServiceException e) {
            log.error(e.getMessage());
            throw e;
        } catch (Exception e) {
            log.error(model.getProvider() + " Streaming Chat 模型配置报错", e);
            return null;
        }
    }

    @Override
    public ChatLanguageModel buildChatLanguageModel(ModelResp model) {
        try {
            if (!whetherCurrentModel(model)) {
                return null;
            }
            if (!basicCheck(model)) {
                return null;
            }
            return OpenAiChatModel.builder()
                .apiKey(model.getApiKey())
                .baseUrl(model.getBaseUrl())
                .modelName(model.getModel())
                .maxTokens(model.getResponseLimit())
                .temperature(model.getTemperature())
                .logRequests(true)
                .logResponses(true)
                .topP(model.getTopP())
                .timeout(Duration.ofMinutes(10))
                .build();
        } catch (ServiceException e) {
            log.error(e.getMessage());
            throw e;
        } catch (Exception e) {
            log.error(model.getProvider() + " Chat 模型配置报错", e);
            return null;
        }
    }

    @Override
    public EmbeddingModel buildEmbedding(ModelResp model) {
        try {
            if (!whetherCurrentModel(model)) {
                return null;
            }
            if (!basicCheck(model)) {
                return null;
            }
            OpenAiEmbeddingModel openAiEmbeddingModel = OpenAiEmbeddingModel.builder()
                .apiKey(model.getApiKey())
                .baseUrl(model.getBaseUrl())
                .modelName(model.getModel())
                .dimensions(model.getDimension())
                .logRequests(true)
                .logResponses(true)
                .dimensions(1024)
                .timeout(Duration.ofMinutes(10))
                .build();
            return openAiEmbeddingModel;
        } catch (ServiceException e) {
            log.error(e.getMessage());
            throw e;
        } catch (Exception e) {
            log.error(model.getProvider() + " Embedding 模型配置报错", e);
            return null;
        }
    }

    @Override
    public ImageModel buildImage(ModelResp model) {
        try {
            if (!whetherCurrentModel(model)) {
                return null;
            }
            if (!basicCheck(model)) {
                return null;
            }
            return OpenAiImageModel.builder()
                .apiKey(model.getApiKey())
                .baseUrl(model.getBaseUrl())
                .modelName(model.getModel())
                .size(model.getImageSize())
                .quality(model.getImageQuality())
                .style(model.getImageStyle())
                .logRequests(true)
                .logResponses(true)
                .timeout(Duration.ofMinutes(10))
                .build();
        } catch (ServiceException e) {
            log.error(e.getMessage());
            throw e;
        } catch (Exception e) {
            log.error(model.getProvider() + " Image 模型配置报错", e);
            return null;
        }

    }
}
