package at.procon.dip.search.rank;

import at.procon.dip.search.api.SearchExecutionContext;
import at.procon.dip.search.api.SearchExecutionPlan;
import at.procon.dip.domain.document.RepresentationType;
import at.procon.dip.search.dto.SearchEngineType;
import at.procon.dip.search.dto.SearchHit;
import at.procon.dip.search.dto.SearchResponse;
import at.procon.dip.search.dto.SearchSortMode;
import at.procon.dip.search.config.DipSearchProperties;
import at.procon.dip.runtime.condition.ConditionalOnRuntimeMode;
import at.procon.dip.runtime.config.RuntimeMode;
import java.util.ArrayList;
import java.util.Comparator;
import java.util.EnumMap;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.UUID;
import lombok.RequiredArgsConstructor;
import org.springframework.stereotype.Component;

@Component
@ConditionalOnRuntimeMode(RuntimeMode.NEW)
@RequiredArgsConstructor
public class DefaultSearchResultFusionService implements SearchResultFusionService {

    private final SearchScoreNormalizer normalizer;
    private final DipSearchProperties properties;

    @Override
    public SearchResponse fuse(SearchExecutionContext context,
                               SearchExecutionPlan plan,
                               Map<SearchEngineType, List<SearchHit>> engineResults) {
        Map<SearchEngineType, List<SearchHit>> normalized = new EnumMap<>(SearchEngineType.class);
        engineResults.forEach((engine, hits) -> normalized.put(engine, normalizer.normalize(engine, hits)));

        List<SearchHit> fused = plan.isCollapseByDocument()
                ? collapse(normalized)
                : mergeWithoutCollapse(normalized);

        sort(fused, plan.getSortMode());
        long total = fused.size();
        int fromIndex = Math.min(context.getPage() * context.getSize(), fused.size());
        int toIndex = Math.min(fromIndex + context.getSize(), fused.size());
        List<SearchHit> paged = fromIndex >= toIndex ? List.of() : fused.subList(fromIndex, toIndex);

        return SearchResponse.builder()
                .hits(paged)
                .page(context.getPage())
                .size(context.getSize())
                .totalHits(total)
                .truncated(total > toIndex)
                .enginesUsed(engineResults.keySet())
                .build();
    }

    private List<SearchHit> collapse(Map<SearchEngineType, List<SearchHit>> normalized) {
        Map<UUID, Aggregate> aggregates = new LinkedHashMap<>();
        normalized.forEach((engine, hits) -> {
            for (SearchHit hit : hits) {
                Aggregate aggregate = aggregates.computeIfAbsent(hit.getDocumentId(), id -> new Aggregate());
                SearchHit currentBestForEngine = aggregate.bestByEngine.get(engine);
                if (currentBestForEngine == null
                        || hit.getNormalizedScore() > currentBestForEngine.getNormalizedScore()
                        || (Double.compare(hit.getNormalizedScore(), currentBestForEngine.getNormalizedScore()) == 0
                            && representationPriority(hit) < representationPriority(currentBestForEngine))) {
                    aggregate.bestByEngine.put(engine, hit);
                }
                if (hit.getRepresentationId() != null) {
                    aggregate.representationIds.add(hit.getRepresentationId());
                }
                if (aggregate.representative == null
                        || hit.getNormalizedScore() > aggregate.representative.getNormalizedScore()
                        || (Double.compare(hit.getNormalizedScore(), aggregate.representative.getNormalizedScore()) == 0
                            && representationPriority(hit) < representationPriority(aggregate.representative))) {
                    aggregate.representative = hit;
                }
            }
        });

        List<SearchHit> fused = new ArrayList<>();
        for (Aggregate aggregate : aggregates.values()) {
            SearchHit representative = aggregate.representative;
            double finalScore = weight(SearchEngineType.POSTGRES_FULLTEXT, aggregate) +
                    weight(SearchEngineType.POSTGRES_TRIGRAM, aggregate) +
                    weight(SearchEngineType.PGVECTOR_SEMANTIC, aggregate) +
                    recencyBoost(representative);
            fused.add(representative.toBuilder()
                    .finalScore(finalScore)
                    .matchedRepresentationCount(aggregate.representationIds.size())
                    .build());
        }
        return fused;
    }

    private double weight(SearchEngineType engineType, Aggregate aggregate) {
        SearchHit hit = aggregate.bestByEngine.get(engineType);
        if (hit == null) {
            return 0.0d;
        }
        DipSearchProperties search = properties;
        return switch (engineType) {
            case POSTGRES_FULLTEXT -> hit.getNormalizedScore() * search.getFulltextWeight();
            case POSTGRES_TRIGRAM -> hit.getNormalizedScore() * search.getTrigramWeight();
            case PGVECTOR_SEMANTIC -> hit.getNormalizedScore() * search.getSemanticWeight();
        };
    }

    private List<SearchHit> mergeWithoutCollapse(Map<SearchEngineType, List<SearchHit>> normalized) {
        List<SearchHit> merged = new ArrayList<>();
        normalized.forEach((engine, hits) -> {
            for (SearchHit hit : hits) {
                double finalScore = switch (engine) {
                    case POSTGRES_FULLTEXT -> hit.getNormalizedScore() * properties.getFulltextWeight();
                    case POSTGRES_TRIGRAM -> hit.getNormalizedScore() * properties.getTrigramWeight();
                    case PGVECTOR_SEMANTIC -> hit.getNormalizedScore() * properties.getSemanticWeight();
                };
                merged.add(hit.toBuilder()
                        .finalScore(finalScore + recencyBoost(hit))
                        .matchedRepresentationCount(1)
                        .build());
            }
        });
        return merged;
    }

    private void sort(List<SearchHit> hits, SearchSortMode sortMode) {
        Comparator<SearchHit> comparator = switch (sortMode) {
            case CREATED_AT_DESC -> Comparator.comparing(SearchHit::getCreatedAt,
                    Comparator.nullsLast(Comparator.reverseOrder()));
            case TITLE_ASC -> Comparator.comparing(hit -> hit.getTitle() == null ? "" : hit.getTitle(),
                    String.CASE_INSENSITIVE_ORDER);
            case SCORE_DESC -> Comparator.comparingDouble(SearchHit::getFinalScore).reversed();
        };
        if (sortMode != SearchSortMode.SCORE_DESC) {
            comparator = comparator.thenComparing(Comparator.comparingDouble(SearchHit::getFinalScore).reversed());
        }
        hits.sort(comparator);
    }

    private double recencyBoost(SearchHit hit) {
        if (properties.getRecencyBoostWeight() <= 0.0d || hit.getCreatedAt() == null) {
            return 0.0d;
        }
        double halfLifeDays = Math.max(1.0d, properties.getRecencyHalfLifeDays());
        double ageDays = Math.max(0.0d, java.time.Duration.between(hit.getCreatedAt(), java.time.OffsetDateTime.now()).toSeconds() / 86400.0d);
        double normalized = Math.exp(-Math.log(2.0d) * (ageDays / halfLifeDays));
        return normalized * properties.getRecencyBoostWeight();
    }

    private int representationPriority(SearchHit hit) {
        if (hit == null) {
            return Integer.MAX_VALUE;
        }
        if (hit.isPrimaryRepresentation()) {
            return 0;
        }
        RepresentationType type = hit.getRepresentationType();
        if (type == RepresentationType.SEMANTIC_TEXT) {
            return 1;
        }
        if (type == RepresentationType.TITLE_ABSTRACT) {
            return 2;
        }
        if (type == RepresentationType.SUMMARY) {
            return 3;
        }
        if (type == RepresentationType.CHUNK) {
            return 4;
        }
        return 5;
    }

    private static final class Aggregate {
        private final Map<SearchEngineType, SearchHit> bestByEngine = new EnumMap<>(SearchEngineType.class);
        private final Set<UUID> representationIds = new java.util.LinkedHashSet<>();
        private SearchHit representative;
    }
}
