package de.frigosped.dc.ai;

import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.mistralai.MistralAiChatModel;
import dev.langchain4j.model.ollama.OllamaChatModel;
import jakarta.enterprise.context.ApplicationScoped;
import jakarta.enterprise.inject.Produces;
import jakarta.inject.Named;
import org.eclipse.microprofile.config.inject.ConfigProperty;

import java.time.Duration;
import java.util.Map;

@ApplicationScoped
public class AiModelProducer {

    @ConfigProperty(name = "dc.ai.provider", defaultValue = "ollama")
    String provider;

    // --- Ollama ---
    @ConfigProperty(name = "dc.ai.ollama.base-url", defaultValue = "")
    String ollamaBaseUrl;

    @ConfigProperty(name = "dc.ai.ollama.api-key", defaultValue = "")
    String ollamaApiKey;

    @ConfigProperty(name = "dc.ai.ollama.main.model", defaultValue = "")
    String ollamaMainModel;

    @ConfigProperty(name = "dc.ai.ollama.catalog.model", defaultValue = "")
    String ollamaCatalogModel;

    @ConfigProperty(name = "dc.ai.ollama.ocr.model", defaultValue = "")
    String ollamaOcrModel;

    // --- Mistral ---
    @ConfigProperty(name = "dc.ai.mistral.api-key", defaultValue = "")
    String mistralApiKey;

    @ConfigProperty(name = "dc.ai.mistral.main.model", defaultValue = "mistral-small-latest")
    String mistralMainModel;

    @ConfigProperty(name = "dc.ai.mistral.catalog.model", defaultValue = "mistral-small-latest")
    String mistralCatalogModel;

    @ConfigProperty(name = "dc.ai.mistral.ocr.model", defaultValue = "pixtral-12b-2409")
    String mistralOcrModel;

    // --- Gemeinsame Parameter ---
    @ConfigProperty(name = "dc.ai.main.timeout", defaultValue = "300s")
    String mainTimeout;

    @ConfigProperty(name = "dc.ai.main.max-retries", defaultValue = "2")
    int mainMaxRetries;

    @ConfigProperty(name = "dc.ai.main.temperature", defaultValue = "0.1")
    double mainTemperature;

    @ConfigProperty(name = "dc.ai.catalog.timeout", defaultValue = "1200s")
    String catalogTimeout;

    @ConfigProperty(name = "dc.ai.catalog.max-retries", defaultValue = "1")
    int catalogMaxRetries;

    @ConfigProperty(name = "dc.ai.catalog.temperature", defaultValue = "0.1")
    double catalogTemperature;

    @ConfigProperty(name = "dc.ai.ocr.timeout", defaultValue = "120s")
    String ocrTimeout;

    @ConfigProperty(name = "dc.ai.ocr.max-retries", defaultValue = "2")
    int ocrMaxRetries;

    @ConfigProperty(name = "dc.ai.log-requests", defaultValue = "false")
    boolean logRequests;

    @ConfigProperty(name = "dc.ai.log-responses", defaultValue = "false")
    boolean logResponses;

    @Produces
    @ApplicationScoped
    @Named("main")
    public ChatLanguageModel mainModel() {
        if (isMistral()) {
            return MistralAiChatModel.builder()
                    .apiKey(mistralApiKey)
                    .modelName(mistralMainModel)
                    .temperature(mainTemperature)
                    .timeout(parseDuration(mainTimeout))
                    .maxRetries(mainMaxRetries)
                    .logRequests(logRequests)
                    .logResponses(logResponses)
                    .build();
        }
        return OllamaChatModel.builder()
                .baseUrl(ollamaBaseUrl)
                .modelName(ollamaMainModel)
                .temperature(mainTemperature)
                .timeout(parseDuration(mainTimeout))
                .maxRetries(mainMaxRetries)
                .customHeaders(Map.of("X-API-KEY", ollamaApiKey))
                .logRequests(logRequests)
                .logResponses(logResponses)
                .build();
    }

    @Produces
    @ApplicationScoped
    @Named("catalog")
    public ChatLanguageModel catalogModel() {
        if (isMistral()) {
            return MistralAiChatModel.builder()
                    .apiKey(mistralApiKey)
                    .modelName(mistralCatalogModel)
                    .temperature(catalogTemperature)
                    .timeout(parseDuration(catalogTimeout))
                    .maxRetries(catalogMaxRetries)
                    .logRequests(logRequests)
                    .logResponses(logResponses)
                    .build();
        }
        return OllamaChatModel.builder()
                .baseUrl(ollamaBaseUrl)
                .modelName(ollamaCatalogModel)
                .temperature(catalogTemperature)
                .timeout(parseDuration(catalogTimeout))
                .maxRetries(catalogMaxRetries)
                .customHeaders(Map.of("X-API-KEY", ollamaApiKey))
                .logRequests(logRequests)
                .logResponses(logResponses)
                .build();
    }

    @Produces
    @ApplicationScoped
    @Named("ocr")
    public ChatLanguageModel ocrModel() {
        if (isMistral()) {
            return MistralAiChatModel.builder()
                    .apiKey(mistralApiKey)
                    .modelName(mistralOcrModel)
                    .timeout(parseDuration(ocrTimeout))
                    .maxRetries(ocrMaxRetries)
                    .logRequests(logRequests)
                    .logResponses(logResponses)
                    .build();
        }
        return OllamaChatModel.builder()
                .baseUrl(ollamaBaseUrl)
                .modelName(ollamaOcrModel)
                .timeout(parseDuration(ocrTimeout))
                .maxRetries(ocrMaxRetries)
                .customHeaders(Map.of("X-API-KEY", ollamaApiKey))
                .logRequests(logRequests)
                .logResponses(logResponses)
                .build();
    }

    private boolean isMistral() {
        return "mistral".equalsIgnoreCase(provider);
    }

    private Duration parseDuration(String value) {
        if (value.endsWith("s")) {
            return Duration.ofSeconds(Long.parseLong(value.replace("s", "")));
        } else if (value.endsWith("m")) {
            return Duration.ofMinutes(Long.parseLong(value.replace("m", "")));
        }
        return Duration.ofSeconds(Long.parseLong(value));
    }
}
