package io.quarkiverse.langchain4j.openai.runtime;

import static io.quarkiverse.langchain4j.runtime.OptionalUtil.firstOrDefault;

import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.Optional;
import java.util.function.Supplier;

import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.DisabledChatLanguageModel;
import dev.langchain4j.model.chat.DisabledStreamingChatLanguageModel;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.model.embedding.DisabledEmbeddingModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.image.DisabledImageModel;
import dev.langchain4j.model.image.ImageModel;
import dev.langchain4j.model.moderation.DisabledModerationModel;
import dev.langchain4j.model.moderation.ModerationModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.openai.OpenAiEmbeddingModel;
import dev.langchain4j.model.openai.OpenAiModerationModel;
import dev.langchain4j.model.openai.OpenAiStreamingChatModel;
import io.quarkiverse.langchain4j.openai.QuarkusOpenAiClient;
import io.quarkiverse.langchain4j.openai.QuarkusOpenAiImageModel;
import io.quarkiverse.langchain4j.openai.runtime.config.ChatModelConfig;
import io.quarkiverse.langchain4j.openai.runtime.config.EmbeddingModelConfig;
import io.quarkiverse.langchain4j.openai.runtime.config.ImageModelConfig;
import io.quarkiverse.langchain4j.openai.runtime.config.LangChain4jOpenAiConfig;
import io.quarkiverse.langchain4j.openai.runtime.config.ModerationModelConfig;
import io.quarkiverse.langchain4j.runtime.NamedModelUtil;
import io.quarkus.runtime.ShutdownContext;
import io.quarkus.runtime.annotations.Recorder;
import io.smallrye.config.ConfigValidationException;

@Recorder
public class OpenAiRecorder {

    private static final String DUMMY_KEY = "dummy";

    public Supplier<ChatLanguageModel> chatModel(LangChain4jOpenAiConfig runtimeConfig, String modelName) {
        LangChain4jOpenAiConfig.OpenAiConfig openAiConfig = correspondingOpenAiConfig(runtimeConfig, modelName);

        if (openAiConfig.enableIntegration()) {
            String apiKey = openAiConfig.apiKey();
            if (DUMMY_KEY.equals(apiKey)) {
                throw new ConfigValidationException(createApiKeyConfigProblems(modelName));
            }
            ChatModelConfig chatModelConfig = openAiConfig.chatModel();
            var builder = OpenAiChatModel.builder()
                    .baseUrl(openAiConfig.baseUrl())
                    .apiKey(apiKey)
                    .timeout(openAiConfig.timeout())
                    .maxRetries(openAiConfig.maxRetries())
                    .logRequests(firstOrDefault(false, chatModelConfig.logRequests(), openAiConfig.logRequests()))
                    .logResponses(firstOrDefault(false, chatModelConfig.logResponses(), openAiConfig.logResponses()))
                    .modelName(chatModelConfig.modelName())
                    .temperature(chatModelConfig.temperature())
                    .topP(chatModelConfig.topP())
                    .presencePenalty(chatModelConfig.presencePenalty())
                    .frequencyPenalty(chatModelConfig.frequencyPenalty())
                    .responseFormat(chatModelConfig.responseFormat().orElse(null));

            openAiConfig.organizationId().ifPresent(builder::organizationId);

            if (chatModelConfig.maxTokens().isPresent()) {
                builder.maxTokens(chatModelConfig.maxTokens().get());
            }

            return new Supplier<>() {
                @Override
                public ChatLanguageModel get() {
                    return builder.build();
                }
            };
        } else {
            return new Supplier<>() {
                @Override
                public ChatLanguageModel get() {
                    return new DisabledChatLanguageModel();
                }
            };
        }
    }

    public Supplier<StreamingChatLanguageModel> streamingChatModel(LangChain4jOpenAiConfig runtimeConfig, String modelName) {
        LangChain4jOpenAiConfig.OpenAiConfig openAiConfig = correspondingOpenAiConfig(runtimeConfig, modelName);

        if (openAiConfig.enableIntegration()) {
            String apiKey = openAiConfig.apiKey();
            if (DUMMY_KEY.equals(apiKey)) {
                throw new ConfigValidationException(createApiKeyConfigProblems(modelName));
            }
            ChatModelConfig chatModelConfig = openAiConfig.chatModel();
            var builder = OpenAiStreamingChatModel.builder()
                    .baseUrl(openAiConfig.baseUrl())
                    .apiKey(apiKey)
                    .timeout(openAiConfig.timeout())
                    .logRequests(firstOrDefault(false, chatModelConfig.logRequests(), openAiConfig.logRequests()))
                    .logResponses(firstOrDefault(false, chatModelConfig.logResponses(), openAiConfig.logResponses()))
                    .modelName(chatModelConfig.modelName())
                    .temperature(chatModelConfig.temperature())
                    .topP(chatModelConfig.topP())
                    .presencePenalty(chatModelConfig.presencePenalty())
                    .frequencyPenalty(chatModelConfig.frequencyPenalty())
                    .responseFormat(chatModelConfig.responseFormat().orElse(null));

            openAiConfig.organizationId().ifPresent(builder::organizationId);

            if (chatModelConfig.maxTokens().isPresent()) {
                builder.maxTokens(chatModelConfig.maxTokens().get());
            }

            return new Supplier<>() {
                @Override
                public StreamingChatLanguageModel get() {
                    return builder.build();
                }
            };
        } else {
            return new Supplier<>() {
                @Override
                public StreamingChatLanguageModel get() {
                    return new DisabledStreamingChatLanguageModel();
                }
            };
        }
    }

    public Supplier<EmbeddingModel> embeddingModel(LangChain4jOpenAiConfig runtimeConfig, String modelName) {
        LangChain4jOpenAiConfig.OpenAiConfig openAiConfig = correspondingOpenAiConfig(runtimeConfig, modelName);

        if (openAiConfig.enableIntegration()) {
            String apiKeyOpt = openAiConfig.apiKey();
            if (DUMMY_KEY.equals(apiKeyOpt)) {
                throw new ConfigValidationException(createApiKeyConfigProblems(modelName));
            }
            EmbeddingModelConfig embeddingModelConfig = openAiConfig.embeddingModel();
            var builder = OpenAiEmbeddingModel.builder()
                    .baseUrl(openAiConfig.baseUrl())
                    .apiKey(apiKeyOpt)
                    .timeout(openAiConfig.timeout())
                    .maxRetries(openAiConfig.maxRetries())
                    .logRequests(firstOrDefault(false, embeddingModelConfig.logRequests(), openAiConfig.logRequests()))
                    .logResponses(firstOrDefault(false, embeddingModelConfig.logResponses(), openAiConfig.logResponses()))
                    .modelName(embeddingModelConfig.modelName());

            if (embeddingModelConfig.user().isPresent()) {
                builder.user(embeddingModelConfig.user().get());
            }

            openAiConfig.organizationId().ifPresent(builder::organizationId);

            return new Supplier<>() {
                @Override
                public EmbeddingModel get() {
                    return builder.build();
                }
            };
        } else {
            return new Supplier<>() {
                @Override
                public EmbeddingModel get() {
                    return new DisabledEmbeddingModel();
                }
            };
        }
    }

    public Supplier<ModerationModel> moderationModel(LangChain4jOpenAiConfig runtimeConfig, String modelName) {
        LangChain4jOpenAiConfig.OpenAiConfig openAiConfig = correspondingOpenAiConfig(runtimeConfig, modelName);

        if (openAiConfig.enableIntegration()) {
            String apiKey = openAiConfig.apiKey();
            if (DUMMY_KEY.equals(apiKey)) {
                throw new ConfigValidationException(createApiKeyConfigProblems(modelName));
            }
            ModerationModelConfig moderationModelConfig = openAiConfig.moderationModel();
            var builder = OpenAiModerationModel.builder()
                    .baseUrl(openAiConfig.baseUrl())
                    .apiKey(apiKey)
                    .timeout(openAiConfig.timeout())
                    .maxRetries(openAiConfig.maxRetries())
                    .logRequests(firstOrDefault(false, moderationModelConfig.logRequests(), openAiConfig.logRequests()))
                    .logResponses(firstOrDefault(false, moderationModelConfig.logResponses(), openAiConfig.logResponses()))
                    .modelName(moderationModelConfig.modelName());

            openAiConfig.organizationId().ifPresent(builder::organizationId);

            return new Supplier<>() {
                @Override
                public ModerationModel get() {
                    return builder.build();
                }
            };
        } else {
            return new Supplier<>() {
                @Override
                public ModerationModel get() {
                    return new DisabledModerationModel();
                }
            };
        }
    }

    public Supplier<ImageModel> imageModel(LangChain4jOpenAiConfig runtimeConfig, String modelName) {
        LangChain4jOpenAiConfig.OpenAiConfig openAiConfig = correspondingOpenAiConfig(runtimeConfig, modelName);

        if (openAiConfig.enableIntegration()) {
            String apiKey = openAiConfig.apiKey();
            if (DUMMY_KEY.equals(apiKey)) {
                throw new ConfigValidationException(createApiKeyConfigProblems(modelName));
            }
            ImageModelConfig imageModelConfig = openAiConfig.imageModel();
            var builder = QuarkusOpenAiImageModel.builder()
                    .baseUrl(openAiConfig.baseUrl())
                    .apiKey(apiKey)
                    .timeout(openAiConfig.timeout())
                    .maxRetries(openAiConfig.maxRetries())
                    .logRequests(firstOrDefault(false, imageModelConfig.logRequests(), openAiConfig.logRequests()))
                    .logResponses(firstOrDefault(false, imageModelConfig.logResponses(), openAiConfig.logResponses()))
                    .modelName(imageModelConfig.modelName())
                    .size(imageModelConfig.size())
                    .quality(imageModelConfig.quality())
                    .style(imageModelConfig.style())
                    .responseFormat(imageModelConfig.responseFormat())
                    .user(imageModelConfig.user());

            openAiConfig.organizationId().ifPresent(builder::organizationId);

            // we persist if the directory was set explicitly and the boolean flag was not set to false
            // or if the boolean flag was set explicitly to true
            Optional<Path> persistDirectory = Optional.empty();
            if (imageModelConfig.persist().isPresent()) {
                if (imageModelConfig.persist().get()) {
                    persistDirectory = imageModelConfig.persistDirectory().or(new Supplier<>() {
                        @Override
                        public Optional<? extends Path> get() {
                            return Optional.of(Paths.get(System.getProperty("java.io.tmpdir"), "dall-e-images"));
                        }
                    });
                }
            } else {
                if (imageModelConfig.persistDirectory().isPresent()) {
                    persistDirectory = imageModelConfig.persistDirectory();
                }
            }

            builder.persistDirectory(persistDirectory);

            return new Supplier<>() {
                @Override
                public ImageModel get() {
                    return builder.build();
                }
            };
        } else {
            return new Supplier<>() {
                @Override
                public ImageModel get() {
                    return new DisabledImageModel();
                }
            };
        }
    }

    private LangChain4jOpenAiConfig.OpenAiConfig correspondingOpenAiConfig(LangChain4jOpenAiConfig runtimeConfig,
            String modelName) {
        LangChain4jOpenAiConfig.OpenAiConfig openAiConfig;
        if (NamedModelUtil.isDefault(modelName)) {
            openAiConfig = runtimeConfig.defaultConfig();
        } else {
            openAiConfig = runtimeConfig.namedConfig().get(modelName);
        }
        return openAiConfig;
    }

    private ConfigValidationException.Problem[] createApiKeyConfigProblems(String modelName) {
        return createConfigProblems("api-key", modelName);
    }

    private ConfigValidationException.Problem[] createConfigProblems(String key, String modelName) {
        return new ConfigValidationException.Problem[] { createConfigProblem(key, modelName) };
    }

    private ConfigValidationException.Problem createConfigProblem(String key, String modelName) {
        return new ConfigValidationException.Problem(String.format(
                "SRCFG00014: The config property quarkus.langchain4j.openai%s%s is required but it could not be found in any config source",
                NamedModelUtil.isDefault(modelName) ? "." : ("." + modelName + "."), key));
    }

    public void cleanUp(ShutdownContext shutdown) {
        shutdown.addShutdownTask(new Runnable() {
            @Override
            public void run() {
                QuarkusOpenAiClient.clearCache();
            }
        });
    }
}

dev.langchain4j.model.openai.OpenAiEmbeddingModel.builder dev.langchain4j.model.openai.OpenAiStreamingChatModel.builder dev.langchain4j.model.openai.OpenAiModerationModel.builder dev.langchain4j.model.openai.OpenAiChatModel.builder
[((2450, 3312), 'dev.langchain4j.model.openai.OpenAiChatModel.builder'), ((2450, 3229), 'dev.langchain4j.model.openai.OpenAiChatModel.builder'), ((2450, 3155), 'dev.langchain4j.model.openai.OpenAiChatModel.builder'), ((2450, 3083), 'dev.langchain4j.model.openai.OpenAiChatModel.builder'), ((2450, 3033), 'dev.langchain4j.model.openai.OpenAiChatModel.builder'), ((2450, 2969), 'dev.langchain4j.model.openai.OpenAiChatModel.builder'), ((2450, 2909), 'dev.langchain4j.model.openai.OpenAiChatModel.builder'), ((2450, 2791), 'dev.langchain4j.model.openai.OpenAiChatModel.builder'), ((2450, 2676), 'dev.langchain4j.model.openai.OpenAiChatModel.builder'), ((2450, 2617), 'dev.langchain4j.model.openai.OpenAiChatModel.builder'), ((2450, 2564), 'dev.langchain4j.model.openai.OpenAiChatModel.builder'), ((2450, 2528), 'dev.langchain4j.model.openai.OpenAiChatModel.builder'), ((4555, 5367), 'dev.langchain4j.model.openai.OpenAiStreamingChatModel.builder'), ((4555, 5284), 'dev.langchain4j.model.openai.OpenAiStreamingChatModel.builder'), ((4555, 5210), 'dev.langchain4j.model.openai.OpenAiStreamingChatModel.builder'), ((4555, 5138), 'dev.langchain4j.model.openai.OpenAiStreamingChatModel.builder'), ((4555, 5088), 'dev.langchain4j.model.openai.OpenAiStreamingChatModel.builder'), ((4555, 5024), 'dev.langchain4j.model.openai.OpenAiStreamingChatModel.builder'), ((4555, 4964), 'dev.langchain4j.model.openai.OpenAiStreamingChatModel.builder'), ((4555, 4846), 'dev.langchain4j.model.openai.OpenAiStreamingChatModel.builder'), ((4555, 4731), 'dev.langchain4j.model.openai.OpenAiStreamingChatModel.builder'), ((4555, 4678), 'dev.langchain4j.model.openai.OpenAiStreamingChatModel.builder'), ((4555, 4642), 'dev.langchain4j.model.openai.OpenAiStreamingChatModel.builder'), ((6642, 7184), 'dev.langchain4j.model.openai.OpenAiEmbeddingModel.builder'), ((6642, 7119), 'dev.langchain4j.model.openai.OpenAiEmbeddingModel.builder'), ((6642, 6996), 'dev.langchain4j.model.openai.OpenAiEmbeddingModel.builder'), ((6642, 6876), 'dev.langchain4j.model.openai.OpenAiEmbeddingModel.builder'), ((6642, 6817), 'dev.langchain4j.model.openai.OpenAiEmbeddingModel.builder'), ((6642, 6764), 'dev.langchain4j.model.openai.OpenAiEmbeddingModel.builder'), ((6642, 6725), 'dev.langchain4j.model.openai.OpenAiEmbeddingModel.builder'), ((8417, 8960), 'dev.langchain4j.model.openai.OpenAiModerationModel.builder'), ((8417, 8894), 'dev.langchain4j.model.openai.OpenAiModerationModel.builder'), ((8417, 8770), 'dev.langchain4j.model.openai.OpenAiModerationModel.builder'), ((8417, 8649), 'dev.langchain4j.model.openai.OpenAiModerationModel.builder'), ((8417, 8590), 'dev.langchain4j.model.openai.OpenAiModerationModel.builder'), ((8417, 8537), 'dev.langchain4j.model.openai.OpenAiModerationModel.builder'), ((8417, 8501), 'dev.langchain4j.model.openai.OpenAiModerationModel.builder'), ((10032, 10845), 'io.quarkiverse.langchain4j.openai.QuarkusOpenAiImageModel.builder'), ((10032, 10794), 'io.quarkiverse.langchain4j.openai.QuarkusOpenAiImageModel.builder'), ((10032, 10723), 'io.quarkiverse.langchain4j.openai.QuarkusOpenAiImageModel.builder'), ((10032, 10670), 'io.quarkiverse.langchain4j.openai.QuarkusOpenAiImageModel.builder'), ((10032, 10613), 'io.quarkiverse.langchain4j.openai.QuarkusOpenAiImageModel.builder'), ((10032, 10562), 'io.quarkiverse.langchain4j.openai.QuarkusOpenAiImageModel.builder'), ((10032, 10501), 'io.quarkiverse.langchain4j.openai.QuarkusOpenAiImageModel.builder'), ((10032, 10382), 'io.quarkiverse.langchain4j.openai.QuarkusOpenAiImageModel.builder'), ((10032, 10266), 'io.quarkiverse.langchain4j.openai.QuarkusOpenAiImageModel.builder'), ((10032, 10207), 'io.quarkiverse.langchain4j.openai.QuarkusOpenAiImageModel.builder'), ((10032, 10154), 'io.quarkiverse.langchain4j.openai.QuarkusOpenAiImageModel.builder'), ((10032, 10118), 'io.quarkiverse.langchain4j.openai.QuarkusOpenAiImageModel.builder')]