package at.procon.dip.embedding.registry;

import at.procon.dip.embedding.config.EmbeddingProperties;
import at.procon.dip.embedding.model.EmbeddingModelDescriptor;
import java.util.List;
import java.util.Optional;
import lombok.RequiredArgsConstructor;
import org.springframework.stereotype.Component;

@Component
@RequiredArgsConstructor
public class EmbeddingModelRegistry {

    private final EmbeddingProperties properties;

    public EmbeddingModelDescriptor getRequired(String modelKey) {
        return find(modelKey)
                .orElseThrow(() -> new IllegalArgumentException("Unknown embedding model key: " + modelKey));
    }

    public Optional<EmbeddingModelDescriptor> find(String modelKey) {
        EmbeddingProperties.ModelProperties model = properties.getModels().get(modelKey);
        if (model == null) {
            return Optional.empty();
        }
        return Optional.of(toDescriptor(modelKey, model));
    }

    public List<EmbeddingModelDescriptor> getActiveModels() {
        return properties.getModels().entrySet().stream()
                .filter(entry -> entry.getValue().isActive())
                .map(entry -> toDescriptor(entry.getKey(), entry.getValue()))
                .toList();
    }

    public String getRequiredDefaultQueryModelKey() {
        if (properties.getDefaultQueryModel() == null || properties.getDefaultQueryModel().isBlank()) {
            throw new IllegalStateException("dip.embedding.default-query-model is not configured");
        }
        return properties.getDefaultQueryModel();
    }

    public String getRequiredDefaultDocumentModelKey() {
        if (properties.getDefaultDocumentModel() == null || properties.getDefaultDocumentModel().isBlank()) {
            throw new IllegalStateException("dip.embedding.default-document-model is not configured");
        }
        return properties.getDefaultDocumentModel();
    }

    private EmbeddingModelDescriptor toDescriptor(String modelKey, EmbeddingProperties.ModelProperties model) {
        return new EmbeddingModelDescriptor(
                modelKey,
                model.getProviderConfigKey(),
                model.getProviderModelKey() == null || model.getProviderModelKey().isBlank() ? modelKey : model.getProviderModelKey(),
                model.getDimensions() == null ? 0 : model.getDimensions(),
                model.getDistanceMetric(),
                model.isSupportsQueryEmbeddingMode(),
                model.isSupportsBatch(),
                model.getMaxInputChars(),
                model.isActive(),
                model.getPrefixMode(),
                model.getQueryPrefix(),
                model.getDocumentPrefix()
        );
    }
}
