package at.procon.dip.embedding.policy;

import at.procon.dip.domain.document.RepresentationType;
import at.procon.dip.domain.document.entity.DocumentTextRepresentation;
import at.procon.dip.domain.document.repository.DocumentTextRepresentationRepository;
import at.procon.dip.embedding.config.EmbeddingProperties;
import at.procon.dip.embedding.model.EmbeddingModelDescriptor;
import java.util.ArrayList;
import java.util.Comparator;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Set;
import java.util.UUID;
import lombok.RequiredArgsConstructor;
import org.springframework.stereotype.Component;

@Component
@RequiredArgsConstructor
public class DefaultEmbeddingSelectionPolicy implements EmbeddingSelectionPolicy {

    private final DocumentTextRepresentationRepository representationRepository;
    private final EmbeddingProperties embeddingProperties;

    @Override
    public List<DocumentTextRepresentation> selectRepresentations(UUID documentId, EmbeddingModelDescriptor model) {
        return selectRepresentations(documentId, model, null);
    }

    @Override
    public List<DocumentTextRepresentation> selectRepresentations(UUID documentId, EmbeddingModelDescriptor model, EmbeddingProfile profile) {
        List<DocumentTextRepresentation> representations = representationRepository.findByDocument_Id(documentId);
        List<DocumentTextRepresentation> selected = new ArrayList<>();
        EmbeddingProperties.IndexingProperties indexing = embeddingProperties.getIndexing();

        for (DocumentTextRepresentation representation : representations) {
            if (include(representation, indexing, profile)) {
                selected.add(representation);
            }
        }

        if (selected.isEmpty()) {
            representationRepository.findFirstByDocument_IdAndPrimaryRepresentationTrue(documentId)
                    .filter(rep -> include(rep, indexing, profile))
                    .ifPresent(selected::add);
        }

        Set<UUID> seen = new LinkedHashSet<>();
        return selected.stream()
                .filter(rep -> seen.add(rep.getId()))
                .sorted(Comparator
                        .comparing(DocumentTextRepresentation::isPrimaryRepresentation).reversed()
                        .thenComparing(rep -> rep.getRepresentationType().ordinal())
                        .thenComparing(rep -> rep.getChunkIndex() == null ? Integer.MAX_VALUE : rep.getChunkIndex()))
                .toList();
    }

    private boolean include(DocumentTextRepresentation representation,
                            EmbeddingProperties.IndexingProperties indexing,
                            EmbeddingProfile profile) {
        if (profile != null && !profile.includes(representation.getRepresentationType())) {
            return false;
        }
        return switch (representation.getRepresentationType()) {
            case SEMANTIC_TEXT -> indexing.isEmbedSemanticText();
            case TITLE_ABSTRACT -> indexing.isEmbedTitleAbstract();
            case CHUNK -> indexing.isEmbedChunks() && effectiveLength(representation) >= indexing.getChunkMinLength();
            case FULLTEXT -> indexing.isEmbedFulltext();
            case SUMMARY -> indexing.isEmbedSummary();
            default -> representation.isPrimaryRepresentation();
        };
    }

    private int effectiveLength(DocumentTextRepresentation representation) {
        if (representation.getCharCount() != null) {
            return representation.getCharCount();
        }
        return representation.getTextBody() == null ? 0 : representation.getTextBody().length();
    }
}
