package at.procon.dip.search.service;

import at.procon.dip.domain.document.entity.DocumentEmbeddingModel;
import at.procon.dip.embedding.config.EmbeddingProperties;
import at.procon.dip.embedding.registry.EmbeddingModelRegistry;
import at.procon.dip.embedding.service.EmbeddingModelCatalogService;
import at.procon.dip.embedding.service.QueryEmbeddingService;
import at.procon.dip.embedding.support.EmbeddingVectorCodec;
import java.util.Optional;
import java.util.UUID;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

@Service
@RequiredArgsConstructor
@Slf4j
public class SemanticQueryEmbeddingService {

    private final EmbeddingProperties embeddingProperties;
    private final EmbeddingModelRegistry embeddingModelRegistry;
    private final EmbeddingModelCatalogService embeddingModelCatalogService;
    private final QueryEmbeddingService queryEmbeddingService;

    public Optional<QueryEmbedding> buildQueryEmbedding(String queryText) {
        return buildQueryEmbedding(queryText, null);
    }

    public Optional<QueryEmbedding> buildQueryEmbedding(String queryText, String requestedModelKey) {
        if (!embeddingProperties.isEnabled()) {
            return Optional.empty();
        }
        if (queryText == null || queryText.isBlank()) {
            return Optional.empty();
        }

        String modelKey = resolveModelKey(requestedModelKey);
        try {
            DocumentEmbeddingModel model = embeddingModelCatalogService.ensureRegistered(modelKey);
            float[] vector = queryEmbeddingService.embedQuery(queryText, modelKey);
            return Optional.of(new QueryEmbedding(
                    model.getId(),
                    modelKey,
                    EmbeddingVectorCodec.toPgVector(vector),
                    vector.length
            ));
        } catch (Exception e) {
            log.warn("Failed to generate semantic query embedding with model {}: {}", modelKey, e.getMessage());
            return Optional.empty();
        }
    }

    private String resolveModelKey(String requestedModelKey) {
        if (requestedModelKey != null && !requestedModelKey.isBlank()) {
            embeddingModelRegistry.getRequired(requestedModelKey);
            return requestedModelKey;
        }
        return embeddingModelRegistry.getRequiredDefaultQueryModelKey();
    }

    public record QueryEmbedding(UUID modelId, String modelKey, String vectorString, int dimensions) {
    }
}
