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

import at.procon.dip.domain.document.DistanceMetric;
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.repository.DocumentEmbeddingModelRepository;
import at.procon.dip.domain.document.repository.DocumentEmbeddingRepository;
import at.procon.dip.domain.document.service.command.RegisterEmbeddingModelCommand;
import java.time.OffsetDateTime;
import java.util.List;
import java.util.Optional;
import java.util.UUID;
import lombok.RequiredArgsConstructor;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

@Service
@RequiredArgsConstructor
@Transactional
public class DocumentEmbeddingService {

    private final DocumentService documentService;
    private final DocumentRepresentationService representationService;
    private final DocumentEmbeddingRepository embeddingRepository;
    private final DocumentEmbeddingModelRepository modelRepository;

    public DocumentEmbeddingModel registerModel(RegisterEmbeddingModelCommand command) {
        DocumentEmbeddingModel model = modelRepository.findByModelKey(command.modelKey())
                .orElseGet(DocumentEmbeddingModel::new);
        model.setModelKey(command.modelKey());
        model.setProvider(command.provider());
        model.setDisplayName(command.displayName());
        model.setDimensions(command.dimensions());
        model.setDistanceMetric(command.distanceMetric() == null ? DistanceMetric.COSINE : command.distanceMetric());
        model.setQueryPrefixRequired(command.queryPrefixRequired());
        model.setActive(command.active());
        return modelRepository.save(model);
    }

    public DocumentEmbedding createPendingEmbedding(UUID documentId, UUID representationId, UUID modelId) {
        DocumentEmbeddingModel model = getRequiredModel(modelId);
        DocumentEmbedding embedding = DocumentEmbedding.builder()
                .document(documentService.getRequired(documentId))
                .representation(representationService.getRequired(representationId))
                .model(model)
                .embeddingDimensions(model.getDimensions())
                .embeddingStatus(EmbeddingStatus.PENDING)
                .build();
        return embeddingRepository.save(embedding);
    }

    public DocumentEmbedding ensurePendingEmbedding(UUID documentId, UUID representationId, UUID modelId) {
        Optional<DocumentEmbedding> existing = embeddingRepository.findByRepresentation_IdAndModel_Id(representationId, modelId);
        if (existing.isPresent()) {
            DocumentEmbedding embedding = existing.get();
            embedding.setDocument(documentService.getRequired(documentId));
            embedding.setRepresentation(representationService.getRequired(representationId));
            embedding.setModel(getRequiredModel(modelId));
            embedding.setEmbeddingDimensions(embedding.getModel().getDimensions());
            embedding.setEmbeddingStatus(EmbeddingStatus.PENDING);
            embedding.setErrorMessage(null);
            embedding.setEmbeddedAt(null);
            return embeddingRepository.save(embedding);
        }
        return createPendingEmbedding(documentId, representationId, modelId);
    }

    public DocumentEmbedding markCompleted(UUID embeddingId, Integer tokenCount) {
        DocumentEmbedding embedding = getRequired(embeddingId);
        embedding.setEmbeddingStatus(EmbeddingStatus.COMPLETED);
        embedding.setTokenCount(tokenCount);
        embedding.setEmbeddedAt(OffsetDateTime.now());
        embedding.setErrorMessage(null);
        return embeddingRepository.save(embedding);
    }

    public DocumentEmbedding markFailed(UUID embeddingId, String errorMessage) {
        DocumentEmbedding embedding = getRequired(embeddingId);
        embedding.setEmbeddingStatus(EmbeddingStatus.FAILED);
        embedding.setErrorMessage(errorMessage);
        embedding.setEmbeddedAt(null);
        return embeddingRepository.save(embedding);
    }

    public DocumentEmbedding markProcessing(UUID embeddingId) {
        DocumentEmbedding embedding = getRequired(embeddingId);
        embedding.setEmbeddingStatus(EmbeddingStatus.PROCESSING);
        embedding.setErrorMessage(null);
        return embeddingRepository.save(embedding);
    }

    public DocumentEmbedding markSkipped(UUID embeddingId, String reason) {
        DocumentEmbedding embedding = getRequired(embeddingId);
        embedding.setEmbeddingStatus(EmbeddingStatus.SKIPPED);
        embedding.setErrorMessage(reason);
        embedding.setEmbeddedAt(OffsetDateTime.now());
        return embeddingRepository.save(embedding);
    }

    @Transactional(readOnly = true)
    public DocumentEmbedding getRequired(UUID embeddingId) {
        return embeddingRepository.findById(embeddingId)
                .orElseThrow(() -> new IllegalArgumentException("Unknown embedding id: " + embeddingId));
    }

    @Transactional(readOnly = true)
    public DocumentEmbeddingModel getRequiredModel(UUID modelId) {
        return modelRepository.findById(modelId)
                .orElseThrow(() -> new IllegalArgumentException("Unknown embedding model id: " + modelId));
    }


    @Transactional(readOnly = true)
    public DocumentEmbeddingModel findActiveModelByKey(String modelKey) {
        return modelRepository.findByModelKey(modelKey)
                .orElseThrow(() -> new IllegalArgumentException("Unknown embedding model key: " + modelKey));
    }

    @Transactional(readOnly = true)
    public List<DocumentEmbedding> findPendingEmbeddings() {
        return embeddingRepository.findByEmbeddingStatus(EmbeddingStatus.PENDING);
    }
}

