package at.procon.dip.domain.document.repository;

import at.procon.dip.domain.document.EmbeddingStatus;
import at.procon.dip.domain.document.entity.DocumentEmbedding;
import java.time.OffsetDateTime;
import java.util.List;
import java.util.Optional;
import java.util.UUID;
import org.springframework.data.domain.Pageable;
import org.springframework.data.jpa.repository.JpaRepository;
import org.springframework.data.jpa.repository.Modifying;
import org.springframework.data.jpa.repository.Query;
import org.springframework.data.repository.query.Param;

public interface DocumentEmbeddingRepository extends JpaRepository<DocumentEmbedding, UUID> {

    List<DocumentEmbedding> findByDocument_Id(UUID documentId);

    List<DocumentEmbedding> findByRepresentation_Id(UUID representationId);

    List<DocumentEmbedding> findByEmbeddingStatus(EmbeddingStatus embeddingStatus);

    Optional<DocumentEmbedding> findByRepresentation_IdAndModel_Id(UUID representationId, UUID modelId);

    @Query("SELECT e.id FROM DocumentEmbedding e WHERE e.embeddingStatus = :status ORDER BY e.createdAt ASC")
    List<UUID> findIdsByEmbeddingStatus(@Param("status") EmbeddingStatus status, Pageable pageable);

    @Query("SELECT e FROM DocumentEmbedding e " +
            "JOIN FETCH e.document d " +
            "JOIN FETCH e.representation r " +
            "JOIN FETCH e.model m " +
            "WHERE e.id = :embeddingId")
    Optional<DocumentEmbedding> findDetailedById(@Param("embeddingId") UUID embeddingId);

    @Modifying
    @Query(value = "UPDATE doc.doc_embedding SET embedding_vector = CAST(:vectorData AS vector), " +
            "embedding_status = 'COMPLETED', embedded_at = CURRENT_TIMESTAMP, updated_at = CURRENT_TIMESTAMP, " +
            "error_message = NULL, token_count = :tokenCount, embedding_dimensions = :dimensions, " +
            "prefix_profile_id = :prefixProfileId WHERE id = :id",
            nativeQuery = true)
    int updateEmbeddingVector(@Param("id") UUID id,
                              @Param("vectorData") float[] vectorData,
                              @Param("tokenCount") Integer tokenCount,
                              @Param("dimensions") Integer dimensions,
                              @Param("prefixProfileId") UUID prefixProfileId);

    @Modifying
    @Query("UPDATE DocumentEmbedding e SET e.embeddingStatus = :status, e.errorMessage = :errorMessage, " +
            "e.embeddedAt = :embeddedAt, e.updatedAt = CURRENT_TIMESTAMP WHERE e.id = :embeddingId")
    int updateEmbeddingStatus(@Param("embeddingId") UUID embeddingId,
                              @Param("status") EmbeddingStatus status,
                              @Param("errorMessage") String errorMessage,
                              @Param("embeddedAt") OffsetDateTime embeddedAt);

    @Query("SELECT e.embeddingStatus, COUNT(e) FROM DocumentEmbedding e GROUP BY e.embeddingStatus")
    List<Object[]> countByEmbeddingStatus();
}
