package ai.decompile.ai.model.repository;

import ai.decompile.ai.model.entity.EmbeddingChunk;
import java.util.List;
import java.util.UUID;
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;
import org.springframework.transaction.annotation.Transactional;

public interface EmbeddingChunkRepository extends JpaRepository<EmbeddingChunk, UUID> {
  @Query(
      value =
          """
        SELECT * FROM embedding_chunk
        WHERE binary_id = :binaryId
          AND 1 - (embedding <=> CAST(:queryEmbedding AS vector)) > :threshold
        ORDER BY embedding <=> CAST(:queryEmbedding AS vector)
        LIMIT :limit
        """,
      nativeQuery = true)
  List<EmbeddingChunk> findSimilar(
      @Param("binaryId") UUID binaryId,
      @Param("queryEmbedding") String queryEmbedding,
      @Param("threshold") double threshold,
      @Param("limit") int limit);

  @Modifying
  @Transactional
  @Query(
      value =
          """
        INSERT INTO embedding_chunk
            (id, binary_id, chunk_type, source_type, source_id, content, embedding, metadata, token_count, created_at)
        VALUES
            (:id, :binaryId, :chunkType, :sourceType, :sourceId, :content, CAST(:embedding AS vector), CAST(:metadata AS jsonb), :tokenCount, :createdAt)
        """,
      nativeQuery = true)
  void insertChunk(
      @Param("id") UUID id,
      @Param("binaryId") UUID binaryId,
      @Param("chunkType") String chunkType,
      @Param("sourceType") String sourceType,
      @Param("sourceId") UUID sourceId,
      @Param("content") String content,
      @Param("embedding") String embedding,
      @Param("metadata") String metadata,
      @Param("tokenCount") Integer tokenCount,
      @Param("createdAt") java.time.Instant createdAt);

  void deleteByBinaryId(UUID binaryId);

  @Query(
      value =
          """
        SELECT * FROM embedding_chunk
        WHERE binary_id = :binaryId
          AND LOWER(content) LIKE LOWER(CONCAT('%', :term, '%'))
        LIMIT :limit
        """,
      nativeQuery = true)
  List<EmbeddingChunk> searchByContent(
      @Param("binaryId") UUID binaryId, @Param("term") String term, @Param("limit") int limit);
}
