package com.arms.egovframework.javaservice.esframework.model.vo;

import com.arms.egovframework.javaservice.esframework.custom.client.CustomDocumentAdapters;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import org.opensearch.search.SearchHit;
import org.opensearch.search.aggregations.Aggregations;
import org.opensearch.search.aggregations.bucket.composite.ParsedComposite;
import org.opensearch.search.aggregations.metrics.TopHits;
import org.springframework.data.elasticsearch.core.AggregationsContainer;
import org.springframework.data.elasticsearch.core.SearchHits;
import org.springframework.data.elasticsearch.core.convert.ElasticsearchConverter;
import org.springframework.data.elasticsearch.core.document.SearchDocument;

import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.stream.Collectors;

@Slf4j
public class CompositeAggregationResultWrapper<T> implements SearchDocResult {

    private final SearchHits<T> searchHits;
    private final List<T> accumulatedDocs;
    @Getter
    private final Map<String, Object> afterKey;
    @Getter
    private final int bucketCount;
    private final Class<T> clazz;
    private final ElasticsearchConverter converter;

    public CompositeAggregationResultWrapper(SearchHits<T> searchHits, ElasticsearchConverter converter, Class<T> clazz) {
        this.searchHits = searchHits;
        this.accumulatedDocs = null;
        this.clazz = clazz;
        this.converter = converter;

        ParsedComposite parsedComposite = extractCompositeAggregation(searchHits.getAggregations());
        this.afterKey = parsedComposite != null ? parsedComposite.afterKey() : null;
        this.bucketCount = parsedComposite != null ? parsedComposite.getBuckets().size() : 0;
    }

    public CompositeAggregationResultWrapper(List<T> accumulatedDocs, Class<T> clazz) {
        this.searchHits = null;
        this.accumulatedDocs = accumulatedDocs;
        this.clazz = clazz;
        this.converter = null;
        this.afterKey = null;
        this.bucketCount = accumulatedDocs.size();
    }

    private T convertToEntity(SearchHit hit, ElasticsearchConverter converter, Class<T> clazz) {
        SearchDocument searchDocument = CustomDocumentAdapters.from(hit);
        return converter.read(clazz, searchDocument);
    }

    private ParsedComposite extractCompositeAggregation(AggregationsContainer<?> aggregationsContainer) {
        return Optional.ofNullable(aggregationsContainer)
                .map(container -> (Aggregations) container.aggregations())
                .map(Aggregations::asList)
                .flatMap(list -> list.stream()
                        .filter(agg -> agg instanceof ParsedComposite)
                        .map(agg -> (ParsedComposite) agg)
                        .findFirst())
                .orElse(null);
    }

    public List<T> toDocs() {
        if (accumulatedDocs != null) {
            return accumulatedDocs;
        }
        return toHitDocs().stream()
                .map(hit -> convertToEntity(hit, converter, clazz))
                .collect(Collectors.toList());
    }

    public List<SearchHit> toHitDocs() {
        if (searchHits == null) {
            return List.of();
        }
        ParsedComposite parsedComposite = extractCompositeAggregation(searchHits.getAggregations());
        if (parsedComposite == null) {
            return List.of();
        }
        return parsedComposite.getBuckets().stream()
                .flatMap(bucket -> bucket.getAggregations().asList().stream())
                .filter(agg -> agg instanceof TopHits)
                .map(agg -> (TopHits) agg)
                .flatMap(topHits -> Arrays.stream(topHits.getHits().getHits()))
                .collect(Collectors.toList());
    }


    public long getTotalHits() {
        if (searchHits == null) {
            return accumulatedDocs != null ? accumulatedDocs.size() : 0;
        }
        return searchHits.getTotalHits();
    }

    public T findFirst() {
        return getT(false);
    }

    public T fetchOnlyOne() {
        return getT(true);
    }

    private T getT(boolean warningVisible) {
        List<T> docs = toDocs();
        if (warningVisible && docs.size() > 1) {
            log.warn("2건 이상이 조회 되었습니다.");
        }

        return docs.stream()
                .findFirst()
                .orElseGet(this::createEmptyInstance);
    }

    private T createEmptyInstance() {
        try {
            return clazz.getConstructor().newInstance();
        } catch (Exception e) {
            throw new IllegalArgumentException(e.getMessage());
        }
    }
}
