package at.procon.dip.embedding.job.repository;

import at.procon.dip.embedding.job.entity.EmbeddingJob;
import at.procon.dip.embedding.model.EmbeddingJobStatus;
import at.procon.dip.embedding.model.EmbeddingJobType;
import java.time.OffsetDateTime;
import java.util.Collection;
import java.util.List;
import java.util.Optional;
import java.util.UUID;
import org.springframework.data.jpa.repository.JpaRepository;
import org.springframework.data.jpa.repository.Query;
import org.springframework.data.repository.query.Param;

public interface EmbeddingJobRepository extends JpaRepository<EmbeddingJob, UUID> {

    Optional<EmbeddingJob> findFirstByRepresentationIdAndModelKeyAndJobTypeAndStatusIn(
            UUID representationId,
            String modelKey,
            EmbeddingJobType jobType,
            Collection<EmbeddingJobStatus> statuses);

    @Query(value = """
            SELECT *
            FROM DOC.doc_embedding_job j
            WHERE j.status IN ('PENDING', 'RETRY_SCHEDULED')
              AND (j.next_retry_at IS NULL OR j.next_retry_at <= :now)
            ORDER BY j.priority DESC, j.updated_at DESC
            FOR UPDATE SKIP LOCKED
            LIMIT :limit
            """, nativeQuery = true)
    List<EmbeddingJob> claimReadyJobs(@Param("now") OffsetDateTime now,
                                      @Param("limit") int limit);

    List<EmbeddingJob> findByDocumentId(UUID documentId);
}
