package at.procon.dip.search.repository;

import at.procon.dip.domain.document.DistanceMetric;
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.dto.SearchMatchField;
import java.util.List;
import java.util.UUID;
import lombok.RequiredArgsConstructor;
import org.springframework.jdbc.core.namedparam.MapSqlParameterSource;
import org.springframework.jdbc.core.namedparam.NamedParameterJdbcTemplate;
import org.springframework.stereotype.Repository;

@Repository
@RequiredArgsConstructor
public class DocumentSemanticSearchRepository {

    private final NamedParameterJdbcTemplate jdbcTemplate;

    public List<SearchHit> search(SearchExecutionContext context,
                                  UUID modelId,
                                  int modelDimensions,
                                  DistanceMetric distanceMetric,
                                  String queryVector,
                                  int limit,
                                  double threshold) {

        if (modelDimensions <= 0) {
            throw new IllegalArgumentException("Semantic search requires a positive model dimension, got: " + modelDimensions);
        }
        if (distanceMetric == null) {
            throw new IllegalArgumentException("Semantic search requires a distance metric");
        }

        String vectorType = "vector(" + modelDimensions + ")";
        String similarityExpr = buildSimilarityExpression(distanceMetric, vectorType);

        StringBuilder sql = new StringBuilder("""
                SELECT
                    d.id AS document_id,
                    dtr.id AS representation_id,
                    CAST(dtr.representation_type AS text) AS representation_type,
                    CAST(d.document_type AS text) AS document_type,
                    CAST(d.document_family AS text) AS document_family,
                    CAST(d.visibility AS text) AS visibility,
                    d.title AS title,
                    d.summary AS summary,
                    COALESCE(dtr.language_code, d.language_code) AS language_code,
                    d.mime_type AS mime_type,
                    dtr.is_primary AS is_primary,
                    dtr.chunk_index AS chunk_index,
                    dtr.chunk_start_offset AS chunk_start_offset,
                    dtr.chunk_end_offset AS chunk_end_offset,
                    d.created_at AS created_at,
                    d.updated_at AS updated_at,
                    LEFT(COALESCE(dtr.text_body, COALESCE(d.summary, d.title, '')), 400) AS snippet,
                """);
        sql.append(similarityExpr).append(" AS score ");
        sql.append("""
                FROM doc.doc_embedding de
                JOIN doc.doc_text_representation dtr ON dtr.id = de.representation_id
                JOIN doc.doc_document d ON d.id = de.document_id
                LEFT JOIN doc.doc_tenant dt ON dt.id = d.owner_tenant_id
                WHERE de.embedding_status = 'COMPLETED'
                  AND de.embedding_vector IS NOT NULL
                  AND de.model_id = :modelId
                  AND de.embedding_dimensions = :modelDimensions
                  AND """);
        sql.append(similarityExpr).append(" >= :threshold ");

        MapSqlParameterSource params = new MapSqlParameterSource();
        params.addValue("queryVector", queryVector);
        params.addValue("modelId", modelId);
        params.addValue("modelDimensions", modelDimensions);
        params.addValue("threshold", threshold);

        SearchSqlFilterSupport.appendCommonFilters(sql, params, context, "d", "dtr", true);
        sql.append(" ORDER BY score DESC, d.updated_at DESC LIMIT :limit");
        params.addValue("limit", limit);

        return jdbcTemplate.query(
                sql.toString(),
                params,
                new SearchHitRowMapper(SearchEngineType.PGVECTOR_SEMANTIC, SearchMatchField.REPRESENTATION_TEXT)
        );
    }

    private String buildSimilarityExpression(DistanceMetric distanceMetric, String vectorType) {
        String lhs = "(de.embedding_vector)::" + vectorType;
        String rhs = "CAST(:queryVector AS " + vectorType + ")";

        return switch (distanceMetric) {
            case COSINE -> "(1 - (" + lhs + " <=> " + rhs + "))";
            case INNER_PRODUCT -> "(-1 * (" + lhs + " <#> " + rhs + "))";
            case EUCLIDEAN, L2 -> "(1 / (1 + (" + lhs + " <-> " + rhs + ")))";
        };
    }
}