package at.procon.dip.embedding.service;

import at.procon.dip.embedding.model.EmbeddingProviderResult;
import at.procon.dip.embedding.model.EmbeddingRequest;
import at.procon.dip.embedding.model.EmbeddingUseCase;
import at.procon.dip.embedding.provider.EmbeddingProvider;
import at.procon.dip.embedding.registry.EmbeddingModelRegistry;
import at.procon.dip.embedding.registry.EmbeddingProviderConfigResolver;
import at.procon.dip.embedding.registry.EmbeddingProviderRegistry;
import java.util.List;
import java.util.Map;
import lombok.RequiredArgsConstructor;
import org.springframework.stereotype.Service;

@Service
@RequiredArgsConstructor
public class EmbeddingExecutionService {

    private final EmbeddingModelRegistry modelRegistry;
    private final EmbeddingProviderConfigResolver providerConfigResolver;
    private final EmbeddingProviderRegistry providerRegistry;

    public EmbeddingProviderResult embedTexts(String modelKey, EmbeddingUseCase useCase, List<String> texts) {
        return embedTexts(modelKey, useCase, texts, Map.of());
    }

    public EmbeddingProviderResult embedTexts(String modelKey,
                                              EmbeddingUseCase useCase,
                                              List<String> texts,
                                              Map<String, Object> providerOptions) {
        var model = modelRegistry.getRequired(modelKey);
        var providerConfig = providerConfigResolver.resolve(model.providerConfigKey());
        EmbeddingProvider provider = providerRegistry.getRequired(providerConfig.providerType());

        if (!provider.supports(model, providerConfig)) {
            throw new IllegalStateException("Provider %s does not support model %s".formatted(
                    provider.providerType(), model.modelKey()));
        }

        EmbeddingRequest request = EmbeddingRequest.builder()
                .modelKey(model.modelKey())
                .useCase(useCase)
                .texts(texts)
                .providerOptions(providerOptions)
                .build();

        return useCase == EmbeddingUseCase.QUERY
                ? provider.embedQuery(providerConfig, model, request)
                : provider.embedDocuments(providerConfig, model, request);
    }
}
