package at.procon.dip.search.engine.semantic;

import at.procon.dip.embedding.config.EmbeddingProperties;
import at.procon.dip.embedding.model.EmbeddingModelDescriptor;
import at.procon.dip.embedding.registry.EmbeddingModelRegistry;
import at.procon.dip.search.api.SearchExecutionContext;
import at.procon.dip.search.dto.SearchEngineType;
import at.procon.dip.search.dto.SearchHit;
import at.procon.dip.search.engine.SearchEngine;
import at.procon.dip.search.repository.DocumentSemanticSearchRepository;
import at.procon.dip.search.service.SemanticQueryEmbeddingService;
import at.procon.dip.search.config.DipSearchProperties;
import at.procon.dip.runtime.condition.ConditionalOnRuntimeMode;
import at.procon.dip.runtime.config.RuntimeMode;
import java.util.List;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

@Component
@ConditionalOnRuntimeMode(RuntimeMode.NEW)
@RequiredArgsConstructor
@Slf4j
public class PgVectorSemanticSearchEngine implements SearchEngine {

    private final EmbeddingProperties embeddingProperties;
    private final EmbeddingModelRegistry embeddingModelRegistry;
    private final DipSearchProperties properties;
    private final SemanticQueryEmbeddingService queryEmbeddingService;
    private final DocumentSemanticSearchRepository repository;

    @Override
    public SearchEngineType type() {
        return SearchEngineType.PGVECTOR_SEMANTIC;
    }

    @Override
    public boolean supports(SearchExecutionContext context) {
        return embeddingProperties.isEnabled()
                && context.getRequest().getQueryText() != null
                && !context.getRequest().getQueryText().isBlank();
    }

    @Override
    public List<SearchHit> execute(SearchExecutionContext context) {
        String requestedModelKey = context.getRequest().getSemanticModelKey();
        EmbeddingModelDescriptor model = resolveModel(requestedModelKey);
        validateModel(model);

        double threshold = context.getRequest().getSemanticSimilarityThreshold() != null
                ? context.getRequest().getSemanticSimilarityThreshold()
                : properties.getSimilarityThreshold();

        return queryEmbeddingService.buildQueryEmbedding(
                        context.getRequest().getQueryText(),
                        model.modelKey())
                .map(query -> repository.search(
                        context,
                        query.modelId(),
                        model.dimensions(),
                        model.distanceMetric(),
                        query.vectorString(),
                        properties.getSemanticCandidateLimit(),
                        threshold))
                .orElseGet(() -> {
                    log.debug("Semantic search skipped because query embedding could not be generated for model {}", model.modelKey());
                    return List.of();
                });
    }

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

    private void validateModel(EmbeddingModelDescriptor model) {
        if (!model.active()) {
            throw new IllegalStateException("Semantic search model is not active: " + model.modelKey());
        }
        if (model.dimensions() <= 0) {
            throw new IllegalStateException("Semantic search model has invalid dimensions: " + model.modelKey() + " -> " + model.dimensions());
        }
        if (model.distanceMetric() == null) {
            throw new IllegalStateException("Semantic search model has no distance metric configured: " + model.modelKey());
        }
        if (!model.supportsQueryEmbeddingMode()) {
            throw new IllegalStateException("Semantic search model does not support query embedding mode: " + model.modelKey());
        }
    }
}