package at.procon.dip.clustering.repository;

import at.procon.dip.clustering.dto.ClusterAssignmentViewResponse;
import at.procon.dip.clustering.dto.ClusterMembersResponse;
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.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;

@Repository
@RequiredArgsConstructor
public class EmbeddingClusterResultQueryRepositoryImpl implements EmbeddingClusterResultQueryRepository {

    private final NamedParameterJdbcTemplate jdbcTemplate;

    @Override
    public List<ClusterAssignmentViewResponse> findAssignments(UUID runId, boolean includeText, int previewLength) {
        String sql = baseSql(includeText) + " where a.cluster_run_id = :runId order by a.cluster_label_raw asc, d.id asc";
        MapSqlParameterSource params = new MapSqlParameterSource()
                .addValue("runId", runId)
                .addValue("previewLength", previewLength);
        return jdbcTemplate.query(sql, params, assignmentRowMapper(includeText));
    }

    @Override
    public List<ClusterMembersResponse> findClusterMembers(UUID runId, UUID clusterId, boolean includeText, int previewLength) {
        String sql = "select c.id as result_cluster_id, c.cluster_label as result_cluster_label, x.* from doc.doc_embedding_cluster c "
                + "join (" + baseSql(includeText) + ") x on x.cluster_id = c.id "
                + "where c.cluster_run_id = :runId and c.id = :clusterId order by x.business_key asc, x.document_id asc";
        MapSqlParameterSource params = new MapSqlParameterSource()
                .addValue("runId", runId)
                .addValue("clusterId", clusterId)
                .addValue("previewLength", previewLength);
        return jdbcTemplate.query(sql, params, (rs, rowNum) -> new ClusterMembersResponse(
                rs.getObject("result_cluster_id", UUID.class),
                rs.getInt("result_cluster_label"),
                assignmentRowMapper(includeText).mapRow(rs, rowNum)
        ));
    }

    private String baseSql(boolean includeText) {
        String textBodyExpression = includeText ? "r.text_body" : "null::text";
        return "select a.id, a.cluster_id, a.embedding_id, a.document_id, a.representation_id, a.cluster_label_raw, "
                + "a.membership_score, a.distance_to_centroid, a.is_noise, d.business_key, d.document_type, "
                + "r.representation_type, r.builder_key, r.language_code, r.char_count as text_length, "
                + "case when r.text_body is null then null "
                + "when char_length(r.text_body) <= :previewLength then r.text_body "
                + "else substring(r.text_body from 1 for :previewLength) end as text_preview, "
                + textBodyExpression + " as text_body "
                + "from doc.doc_embedding_cluster_assignment a "
                + "join doc.doc_document d on d.id = a.document_id "
                + "join doc.doc_text_representation r on r.id = a.representation_id";
    }

    private RowMapper<ClusterAssignmentViewResponse> assignmentRowMapper(boolean includeText) {
        return (rs, rowNum) -> new ClusterAssignmentViewResponse(
                rs.getObject("id", UUID.class),
                rs.getObject("cluster_id", UUID.class),
                rs.getObject("embedding_id", UUID.class),
                rs.getObject("document_id", UUID.class),
                rs.getObject("representation_id", UUID.class),
                rs.getInt("cluster_label_raw"),
                rs.getObject("membership_score", Double.class),
                rs.getObject("distance_to_centroid", Double.class),
                rs.getBoolean("is_noise"),
                rs.getString("business_key"),
                enumValue(DocumentType.class, rs.getString("document_type")),
                enumValue(RepresentationType.class, rs.getString("representation_type")),
                rs.getString("builder_key"),
                rs.getString("language_code"),
                rs.getObject("text_length", Integer.class),
                rs.getString("text_preview"),
                includeText ? rs.getString("text_body") : null
        );
    }

    private <T extends Enum<T>> T enumValue(Class<T> type, String value) {
        return value == null ? null : Enum.valueOf(type, value);
    }
}
