package dev.langchain4j.zhipu.spring;

import dev.langchain4j.model.zhipu.ZhipuAiChatModel;
import dev.langchain4j.model.zhipu.ZhipuAiEmbeddingModel;
import dev.langchain4j.model.zhipu.ZhipuAiStreamingChatModel;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import static dev.langchain4j.zhipu.spring.Properties.PREFIX;

@Configuration
@EnableConfigurationProperties(Properties.class)
public class ZhipuAutoConfig {

    @Bean
    @ConditionalOnProperty(PREFIX + ".chat-model.api-key")
    ZhipuAiChatModel zhipuAiChatModel(Properties properties) {
        ChatModelProperties chatModelProperties = properties.getChatModel();
        return ZhipuAiChatModel.builder().baseUrl(chatModelProperties.getBaseUrl())
                .apiKey(chatModelProperties.getApiKey()).model(chatModelProperties.getModelName())
                .temperature(chatModelProperties.getTemperature())
                .topP(chatModelProperties.getTopP()).maxRetries(chatModelProperties.getMaxRetries())
                .maxToken(chatModelProperties.getMaxToken())
                .logRequests(chatModelProperties.getLogRequests())
                .logResponses(chatModelProperties.getLogResponses()).build();
    }

    @Bean
    @ConditionalOnProperty(PREFIX + ".streaming-chat-model.api-key")
    ZhipuAiStreamingChatModel zhipuStreamingChatModel(Properties properties) {
        ChatModelProperties chatModelProperties = properties.getStreamingChatModel();
        return ZhipuAiStreamingChatModel.builder().baseUrl(chatModelProperties.getBaseUrl())
                .apiKey(chatModelProperties.getApiKey()).model(chatModelProperties.getModelName())
                .temperature(chatModelProperties.getTemperature())
                .topP(chatModelProperties.getTopP()).maxToken(chatModelProperties.getMaxToken())
                .logRequests(chatModelProperties.getLogRequests())
                .logResponses(chatModelProperties.getLogResponses()).build();
    }

    @Bean
    @ConditionalOnProperty(PREFIX + ".embedding-model.api-key")
    ZhipuAiEmbeddingModel zhipuEmbeddingModel(Properties properties) {
        EmbeddingModelProperties embeddingModelProperties = properties.getEmbeddingModel();
        return ZhipuAiEmbeddingModel.builder().baseUrl(embeddingModelProperties.getBaseUrl())
                .apiKey(embeddingModelProperties.getApiKey())
                .model(embeddingModelProperties.getModel())
                .maxRetries(embeddingModelProperties.getMaxRetries())
                .logRequests(embeddingModelProperties.getLogRequests())
                .logResponses(embeddingModelProperties.getLogResponses()).build();
    }
}