package at.procon.dip.clustering.service;

import at.procon.dip.clustering.dto.ClusteringEngineAssignment;
import at.procon.dip.clustering.dto.ClusteringEngineCluster;
import at.procon.dip.clustering.dto.ClusteringEngineResult;
import at.procon.dip.clustering.entity.EmbeddingCluster;
import at.procon.dip.clustering.entity.EmbeddingClusterRun;
import at.procon.dip.clustering.repository.EmbeddingClusterRepository;
import java.time.OffsetDateTime;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
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.Service;
import org.springframework.transaction.annotation.Transactional;

@Service
@RequiredArgsConstructor
public class EmbeddingClusterPersistenceService {

    private final EmbeddingClusterRepository clusterRepository;
    private final NamedParameterJdbcTemplate jdbcTemplate;

    @Transactional
    public void persist(EmbeddingClusterRun run,
                        List<ClusteringEngineCluster> clusters,
                        List<ClusteringEngineAssignment> assignments) {
        persist(run, new ClusteringEngineResult(clusters, assignments));
    }

    @Transactional
    public void persist(EmbeddingClusterRun run, ClusteringEngineResult result) {
        Map<Integer, EmbeddingCluster> clusterByLabel = new HashMap<>();
        for (ClusteringEngineCluster cluster : result.clusters()) {
            EmbeddingCluster saved = clusterRepository.save(EmbeddingCluster.builder()
                    .clusterRun(run)
                    .clusterLabel(cluster.clusterLabel())
                    .itemCount(cluster.itemCount())
                    .noiseCluster(cluster.noiseCluster())
                    .build());
            clusterRepository.flush();
            clusterByLabel.put(cluster.clusterLabel(), saved);
        }

        if (result.assignments().isEmpty()) {
            return;
        }

        MapSqlParameterSource[] batch = result.assignments().stream()
                .map(assignment -> toInsertParams(run, clusterByLabel.get(assignment.clusterLabel()), assignment))
                .toArray(MapSqlParameterSource[]::new);

        jdbcTemplate.batchUpdate("""
                insert into doc.doc_embedding_cluster_assignment (
                    id,
                    cluster_run_id,
                    cluster_id,
                    embedding_id,
                    document_id,
                    representation_id,
                    cluster_label_raw,
                    membership_score,
                    distance_to_centroid,
                    is_noise,
                    created_at
                )
                select
                    :id,
                    :clusterRunId,
                    :clusterId,
                    e.id,
                    e.document_id,
                    e.representation_id,
                    :clusterLabelRaw,
                    :membershipScore,
                    :distanceToCentroid,
                    :isNoise,
                    :createdAt
                from doc.doc_embedding e
                where e.id = :embeddingId
                """, batch);
    }

    private MapSqlParameterSource toInsertParams(
            EmbeddingClusterRun run,
            EmbeddingCluster cluster,
            ClusteringEngineAssignment assignment
    ) {
        return new MapSqlParameterSource()
                .addValue("id", UUID.randomUUID())
                .addValue("clusterRunId", run.getId())
                .addValue("clusterId", cluster == null ? null : cluster.getId())
                .addValue("embeddingId", assignment.embeddingId())
                .addValue("clusterLabelRaw", assignment.clusterLabel())
                .addValue("membershipScore", assignment.membershipScore())
                .addValue("distanceToCentroid", assignment.distanceToCentroid())
                .addValue("isNoise", assignment.noise())
                .addValue("createdAt", OffsetDateTime.now());
    }
}
