package at.procon.dip.embedding.service;

import at.procon.dip.domain.document.EmbeddingStatus;
import at.procon.dip.domain.document.entity.DocumentEmbedding;
import at.procon.dip.domain.document.entity.DocumentEmbeddingModel;
import at.procon.dip.domain.document.entity.DocumentTextRepresentation;
import at.procon.dip.domain.document.repository.DocumentEmbeddingRepository;
import at.procon.dip.domain.document.repository.DocumentTextRepresentationRepository;
import at.procon.dip.domain.document.service.DocumentEmbeddingService;
import at.procon.dip.embedding.model.EmbeddingPrefixMode;
import at.procon.dip.embedding.model.EmbeddingProviderResult;
import java.time.OffsetDateTime;
import java.util.UUID;
import lombok.RequiredArgsConstructor;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

@Service
@RequiredArgsConstructor
@Transactional
public class EmbeddingPersistenceService {

    private final DocumentTextRepresentationRepository representationRepository;
    private final DocumentEmbeddingService documentEmbeddingService;
    private final DocumentEmbeddingRepository embeddingRepository;
    private final EmbeddingModelCatalogService modelCatalogService;
    private final EmbeddingPrefixProfileService embeddingPrefixProfileService;

    public DocumentEmbedding ensurePending(UUID representationId, String modelKey) {
        DocumentTextRepresentation representation = representationRepository.findById(representationId)
                .orElseThrow(() -> new IllegalArgumentException("Unknown representation id: " + representationId));
        DocumentEmbeddingModel model = modelCatalogService.ensureRegistered(modelKey);
        return documentEmbeddingService.ensurePendingEmbedding(
                representation.getDocument().getId(),
                representation.getId(),
                model.getId()
        );
    }

    public DocumentEmbedding markProcessing(UUID embeddingId) {
        return documentEmbeddingService.markProcessing(embeddingId);
    }

    public void saveCompleted(UUID embeddingId, EmbeddingProviderResult result) {
        if (result.vectors() == null || result.vectors().isEmpty()) {
            throw new IllegalArgumentException("Embedding provider result contains no vectors");
        }
        saveCompleted(embeddingId, result.vectors().getFirst(), result.tokenCount(), result.prefixMode(), result.appliedPrefix());
    }

    public void saveCompleted(UUID embeddingId, float[] vector, Integer tokenCount) {
        saveCompleted(embeddingId, vector, tokenCount, EmbeddingPrefixMode.OFF, null);
    }

    public void saveCompleted(UUID embeddingId,
                              float[] vector,
                              Integer tokenCount,
                              EmbeddingPrefixMode prefixMode,
                              String appliedPrefix) {
        if (vector == null || vector.length == 0) {
            throw new IllegalArgumentException("Embedding vector must not be empty");
        }
        UUID prefixProfileId = embeddingPrefixProfileService.resolveProfileId(prefixMode, appliedPrefix);
        embeddingRepository.updateEmbeddingVector(
                embeddingId,
                vector,
                tokenCount,
                vector.length,
                prefixProfileId
        );
    }

    public void markFailed(UUID embeddingId, String errorMessage) {
        embeddingRepository.updateEmbeddingStatus(
                embeddingId,
                EmbeddingStatus.FAILED,
                errorMessage,
                null
        );
    }

    public void markSkipped(UUID embeddingId, String reason) {
        embeddingRepository.updateEmbeddingStatus(
                embeddingId,
                EmbeddingStatus.SKIPPED,
                reason,
                OffsetDateTime.now()
        );
    }
}
