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

import at.procon.dip.domain.document.entity.DocumentTextRepresentation;
import at.procon.dip.embedding.config.EmbeddingProperties;
import at.procon.dip.embedding.job.entity.EmbeddingJob;
import at.procon.dip.embedding.job.repository.EmbeddingJobRepository;
import at.procon.dip.embedding.model.EmbeddingJobStatus;
import at.procon.dip.embedding.model.EmbeddingJobType;
import at.procon.dip.embedding.policy.EmbeddingProfile;
import at.procon.dip.embedding.policy.EmbeddingSelectionPolicy;
import at.procon.dip.embedding.registry.EmbeddingModelRegistry;
import java.time.Duration;
import java.time.OffsetDateTime;
import java.util.List;
import java.util.Set;
import java.util.UUID;
import lombok.RequiredArgsConstructor;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Propagation;
import org.springframework.transaction.annotation.Transactional;

@Service
@RequiredArgsConstructor
@Transactional
public class EmbeddingJobService {

    private static final Set<EmbeddingJobStatus> ACTIVE_STATUSES = Set.of(
            EmbeddingJobStatus.PENDING,
            EmbeddingJobStatus.IN_PROGRESS,
            EmbeddingJobStatus.RETRY_SCHEDULED
    );

    private final EmbeddingJobRepository jobRepository;
    private final EmbeddingSelectionPolicy selectionPolicy;
    private final EmbeddingModelRegistry modelRegistry;
    private final EmbeddingProperties properties;

    public List<EmbeddingJob> enqueueForDocument(UUID documentId) {
        return enqueueForDocument(documentId, modelRegistry.getRequiredDefaultDocumentModelKey());
    }

    public List<EmbeddingJob> enqueueForDocument(UUID documentId, String modelKey) {
        var model = modelRegistry.getRequired(modelKey);
        List<DocumentTextRepresentation> selected = selectionPolicy.selectRepresentations(documentId, model);
        return selected.stream()
                .map(representation -> enqueueForRepresentation(documentId, representation.getId(), modelKey, EmbeddingJobType.DOCUMENT_EMBED))
                .toList();
    }

    public List<EmbeddingJob> enqueueForDocument(UUID documentId, String modelKey, EmbeddingProfile profile) {
        var model = modelRegistry.getRequired(modelKey);
        List<DocumentTextRepresentation> selected = selectionPolicy.selectRepresentations(documentId, model, profile);
        return selected.stream()
                .map(representation -> enqueueForRepresentation(documentId, representation.getId(), modelKey, EmbeddingJobType.DOCUMENT_EMBED))
                .toList();
    }

    public EmbeddingJob enqueueForRepresentation(UUID documentId, UUID representationId, String modelKey, EmbeddingJobType jobType) {
        return jobRepository.findFirstByRepresentationIdAndModelKeyAndJobTypeAndStatusIn(
                        representationId,
                        modelKey,
                        jobType,
                        ACTIVE_STATUSES
                )
                .orElseGet(() -> jobRepository.save(EmbeddingJob.builder()
                        .documentId(documentId)
                        .representationId(representationId)
                        .modelKey(modelKey)
                        .jobType(jobType)
                        .status(EmbeddingJobStatus.PENDING)
                        .priority(0)
                        .attemptCount(0)
                        .build()));
    }

    @Transactional(propagation = Propagation.REQUIRES_NEW)
    public List<EmbeddingJob> claimNextReadyJobs(int limit) {
        List<EmbeddingJob> jobs = jobRepository.claimReadyJobs(OffsetDateTime.now(), limit);
        jobs.forEach(job -> {
            job.setStatus(EmbeddingJobStatus.IN_PROGRESS);
            job.setAttemptCount(job.getAttemptCount() + 1);
            job.setLastError(null);
            job.setNextRetryAt(null);
        });
        return jobRepository.saveAll(jobs);
    }

    public EmbeddingJob getRequired(UUID jobId) {
        return jobRepository.findById(jobId)
                .orElseThrow(() -> new IllegalArgumentException("Unknown embedding job id: " + jobId));
    }

    public void markCompleted(UUID jobId, String providerRequestId) {
        EmbeddingJob job = getRequired(jobId);
        job.setStatus(EmbeddingJobStatus.DONE);
        job.setProviderRequestId(providerRequestId);
        job.setLastError(null);
        job.setNextRetryAt(null);
        jobRepository.save(job);
    }

    public void markFailed(UUID jobId, String errorMessage, boolean retryable) {
        EmbeddingJob job = getRequired(jobId);
        job.setLastError(errorMessage);
        if (retryable && job.getAttemptCount() < properties.getJobs().getMaxRetries()) {
            job.setStatus(EmbeddingJobStatus.RETRY_SCHEDULED);
            job.setNextRetryAt(OffsetDateTime.now().plus(nextDelay(job.getAttemptCount())));
        } else {
            job.setStatus(EmbeddingJobStatus.FAILED);
            job.setNextRetryAt(null);
        }
        jobRepository.save(job);
    }

    private Duration nextDelay(int attemptCount) {
        long factor = Math.max(0, attemptCount - 1);
        Duration candidate = properties.getJobs().getInitialRetryDelay().multipliedBy(1L << Math.min(factor, 10));
        return candidate.compareTo(properties.getJobs().getMaxRetryDelay()) > 0
                ? properties.getJobs().getMaxRetryDelay()
                : candidate;
    }
}
