package at.procon.dip.clustering.repository;

import at.procon.dip.clustering.dto.EmbeddingSelectionSpec;
import at.procon.dip.clustering.dto.SelectedEmbeddingRow;
import at.procon.dip.domain.document.DocumentFamily;
import at.procon.dip.domain.document.DocumentType;
import at.procon.dip.domain.document.RepresentationType;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.List;
import java.util.UUID;
import lombok.RequiredArgsConstructor;
import org.springframework.jdbc.core.RowMapper;
import org.springframework.jdbc.core.namedparam.MapSqlParameterSource;
import org.springframework.jdbc.core.namedparam.NamedParameterJdbcTemplate;
import org.springframework.stereotype.Repository;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;

@Repository
@RequiredArgsConstructor
public class DocumentEmbeddingClusterSelectionRepositoryImpl implements DocumentEmbeddingClusterSelectionRepository {

    private final NamedParameterJdbcTemplate jdbcTemplate;

    @Override
    public List<SelectedEmbeddingRow> findSelection(EmbeddingSelectionSpec spec) {
        StringBuilder sql = new StringBuilder("""
                select
                    e.id as embedding_id,
                    d.id as document_id,
                    r.id as representation_id,
                    e.model_id,
                    e.prefix_profile_id,
                    d.document_type,
                    d.document_family,
                    r.representation_type,
                    r.builder_key,
                    r.language_code,
                    d.business_key,
                    e.embedding_vector::text as embedding_vector_text
                from doc.doc_embedding e
                join doc.doc_document d on d.id = e.document_id
                join doc.doc_text_representation r on r.id = e.representation_id
                where e.embedding_status = 'COMPLETED'
                  and e.embedding_vector is not null
                  and e.prefix_profile_id is not null
                """);
        MapSqlParameterSource params = new MapSqlParameterSource();
        applyFilters(spec, sql, params);
        sql.append(" order by e.created_at asc");

        List<RawSelectedEmbeddingRow> rawRows = jdbcTemplate.query(
                sql.toString(),
                params,
                new RawSelectedEmbeddingRowMapper());

        return rawRows.stream()
                .map(this::toSelectedEmbeddingRow)
                .toList();
    }

    @Override
    public long countSelection(EmbeddingSelectionSpec spec) {
        StringBuilder sql = new StringBuilder("""
                select count(*)
                from doc.doc_embedding e
                join doc.doc_document d on d.id = e.document_id
                join doc.doc_text_representation r on r.id = e.representation_id
                where e.embedding_status = 'COMPLETED'
                  and e.embedding_vector is not null
                  and e.prefix_profile_id is not null
                """);
        MapSqlParameterSource params = new MapSqlParameterSource();
        applyFilters(spec, sql, params);
        Long result = jdbcTemplate.queryForObject(sql.toString(), params, Long.class);
        return result == null ? 0L : result;
    }

    private void applyFilters(EmbeddingSelectionSpec spec, StringBuilder sql, MapSqlParameterSource params) {
        if (spec == null) {
            return;
        }
        if (!CollectionUtils.isEmpty(spec.documentTypes())) {
            sql.append(" and d.document_type in (:documentTypes)");
            params.addValue("documentTypes", enumNames(spec.documentTypes()));
        }
        if (!CollectionUtils.isEmpty(spec.documentFamilies())) {
            sql.append(" and d.document_family in (:documentFamilies)");
            params.addValue("documentFamilies", enumNames(spec.documentFamilies()));
        }
        if (!CollectionUtils.isEmpty(spec.representationTypes())) {
            sql.append(" and r.representation_type in (:representationTypes)");
            params.addValue("representationTypes", enumNames(spec.representationTypes()));
        }
        if (!CollectionUtils.isEmpty(spec.embeddingStatuses())) {
            sql.append(" and e.embedding_status in (:embeddingStatuses)");
            params.addValue("embeddingStatuses", enumNames(spec.embeddingStatuses()));
        }
        if (!CollectionUtils.isEmpty(spec.modelIds())) {
            sql.append(" and e.model_id in (:modelIds)");
            params.addValue("modelIds", spec.modelIds());
        }
        if (!CollectionUtils.isEmpty(spec.prefixProfileIds())) {
            sql.append(" and e.prefix_profile_id in (:prefixProfileIds)");
            params.addValue("prefixProfileIds", spec.prefixProfileIds());
        }
        if (!CollectionUtils.isEmpty(spec.builderKeys())) {
            sql.append(" and r.builder_key in (:builderKeys)");
            params.addValue("builderKeys", spec.builderKeys());
        }
        if (!CollectionUtils.isEmpty(spec.languageCodes())) {
            sql.append(" and r.language_code in (:languageCodes)");
            params.addValue("languageCodes", spec.languageCodes());
        }
        if (!CollectionUtils.isEmpty(spec.ownerTenantIds())) {
            sql.append(" and d.owner_tenant_id in (:ownerTenantIds)");
            params.addValue("ownerTenantIds", spec.ownerTenantIds());
        }
        if (StringUtils.hasText(spec.businessKeyLike())) {
            sql.append(" and d.business_key like :businessKeyLike");
            params.addValue("businessKeyLike", spec.businessKeyLike());
        }
        if (spec.createdFrom() != null) {
            sql.append(" and d.created_at >= :createdFrom");
            params.addValue("createdFrom", spec.createdFrom());
        }
        if (spec.createdTo() != null) {
            sql.append(" and d.created_at < :createdTo");
            params.addValue("createdTo", spec.createdTo());
        }
        if (Boolean.TRUE.equals(spec.primaryRepresentationOnly())) {
            sql.append(" and r.is_primary = true");
        }
    }

    private List<String> enumNames(Iterable<?> values) {
        List<String> result = new ArrayList<>();
        for (Object value : values) {
            result.add(String.valueOf(value));
        }
        return result;
    }

    private SelectedEmbeddingRow toSelectedEmbeddingRow(RawSelectedEmbeddingRow row) {
        return new SelectedEmbeddingRow(
                row.embeddingId(),
                row.documentId(),
                row.representationId(),
                row.modelId(),
                row.prefixProfileId(),
                row.documentType(),
                row.documentFamily(),
                row.representationType(),
                row.builderKey(),
                row.languageCode(),
                row.businessKey(),
                parseVector(row.embeddingVectorText()));
    }

    private float[] parseVector(String raw) {
        if (raw == null) {
            return null;
        }
        String value = raw.trim();
        if (value.length() < 2) {
            return new float[0];
        }
        if (value.charAt(0) == '[' && value.charAt(value.length() - 1) == ']') {
            value = value.substring(1, value.length() - 1);
        }
        if (value.isBlank()) {
            return new float[0];
        }
        String[] parts = value.split(",");
        float[] result = new float[parts.length];
        for (int i = 0; i < parts.length; i++) {
            result[i] = Float.parseFloat(parts[i].trim());
        }
        return result;
    }

    private record RawSelectedEmbeddingRow(
            UUID embeddingId,
            UUID documentId,
            UUID representationId,
            UUID modelId,
            UUID prefixProfileId,
            DocumentType documentType,
            DocumentFamily documentFamily,
            RepresentationType representationType,
            String builderKey,
            String languageCode,
            String businessKey,
            String embeddingVectorText
    ) {
    }

    private static class RawSelectedEmbeddingRowMapper implements RowMapper<RawSelectedEmbeddingRow> {
        @Override
        public RawSelectedEmbeddingRow mapRow(ResultSet rs, int rowNum) throws SQLException {
            return new RawSelectedEmbeddingRow(
                    rs.getObject("embedding_id", UUID.class),
                    rs.getObject("document_id", UUID.class),
                    rs.getObject("representation_id", UUID.class),
                    rs.getObject("model_id", UUID.class),
                    rs.getObject("prefix_profile_id", UUID.class),
                    DocumentType.valueOf(rs.getString("document_type")),
                    DocumentFamily.valueOf(rs.getString("document_family")),
                    RepresentationType.valueOf(rs.getString("representation_type")),
                    rs.getString("builder_key"),
                    rs.getString("language_code"),
                    rs.getString("business_key"),
                    rs.getString("embedding_vector_text")
            );
        }
    }
}
