package org.xbib.elasticsearch.plugin.bundle.query.decompound;

import org.apache.lucene.index.Term;
import org.apache.lucene.search.BooleanClause;
import org.apache.lucene.search.BooleanQuery;
import org.apache.lucene.search.BoostQuery;
import org.apache.lucene.search.ConstantScoreQuery;
import org.apache.lucene.search.DisjunctionMaxQuery;
import org.apache.lucene.search.PhraseQuery;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.spans.SpanNearQuery;
import org.apache.lucene.search.spans.SpanTermQuery;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;

public class QueryTransformer {

    private final List<QueryHandler> queryHandlers = Arrays.asList(new BooleanQueryHandler(),
            new BoostQueryHandler(), new DisjunctionMaxQueryHandler(), new ConstantScoreQueryHandler(),
            new PhraseQueryHandler());

    public Query transform(Query query) {
        for (QueryHandler queryHandler : queryHandlers) {
            if (queryHandler.accept(query)) {
                return queryHandler.handle(query, this);
            }
        }
        return query;
    }

    interface QueryHandler {

        boolean accept(Query query);

        Query handle(Query query, QueryTransformer queryTransformer);
    }

    class BooleanQueryHandler implements QueryHandler {

        @Override
        public boolean accept(Query query) {
            return query instanceof BooleanQuery;
        }

        @Override
        public Query handle(Query query, QueryTransformer queryTransformer) {
            BooleanQuery booleanQuery = (BooleanQuery) query;
            boolean changed = false;
            BooleanQuery.Builder builder = new BooleanQuery.Builder();
            for (BooleanClause clause: booleanQuery.clauses()) {
                Query newClauseQuery = queryTransformer.transform(clause.getQuery());
                if (newClauseQuery != clause.getQuery()) {
                    changed = true;
                    builder.add(new BooleanClause(newClauseQuery, clause.getOccur()));
                } else {
                    builder.add(clause);
                }
            }
            if (changed) {
                builder.setMinimumNumberShouldMatch(booleanQuery.getMinimumNumberShouldMatch());
                return builder.build();
            }
            return query;
        }
    }

    class BoostQueryHandler implements QueryHandler {

        @Override
        public boolean accept(Query query) {
            return query instanceof BoostQuery;
        }

        @Override
        public Query handle(Query query, QueryTransformer queryTransformer) {
            BoostQuery boostQuery = (BoostQuery) query;
            Query newInnerBoostQuery = queryTransformer.transform(boostQuery.getQuery());
            if (newInnerBoostQuery != boostQuery.getQuery()) {
                return new BoostQuery(newInnerBoostQuery, boostQuery.getBoost());
            }
            return query;
        }
    }

    class DisjunctionMaxQueryHandler implements QueryHandler {

        @Override
        public boolean accept(Query query) {
            return query instanceof DisjunctionMaxQuery;
        }

        @Override
        public Query handle(Query query, QueryTransformer queryTransformer) {
            DisjunctionMaxQuery disjunctionMaxQuery = (DisjunctionMaxQuery) query;
            boolean changed = false;
            List<Query> innerQueries = new ArrayList<>();
            for (Query innerQuery: disjunctionMaxQuery.getDisjuncts()) {
                Query newInnerQuery = queryTransformer.transform(innerQuery);
                if (newInnerQuery != innerQuery) {
                    changed = true;
                    innerQueries.add(newInnerQuery);
                } else {
                    innerQueries.add(innerQuery);
                }
            }
            if (changed) {
                return new DisjunctionMaxQuery(innerQueries, disjunctionMaxQuery.getTieBreakerMultiplier());
            }
            return query;
        }
    }

    class ConstantScoreQueryHandler implements QueryHandler {

        @Override
        public boolean accept(Query query) {
            return query instanceof ConstantScoreQuery;
        }

        @Override
        public Query handle(Query query, QueryTransformer queryTransformer) {
            ConstantScoreQuery constantScoreQuery = (ConstantScoreQuery) query;
            Query newInnerConstantScoreQuery = queryTransformer.transform(constantScoreQuery.getQuery());
            if (newInnerConstantScoreQuery != constantScoreQuery.getQuery()) {
                return new ConstantScoreQuery(newInnerConstantScoreQuery);
            }
            return query;
        }
    }

    class PhraseQueryHandler implements QueryHandler {

        @Override
        public boolean accept(Query query) {
            return query instanceof PhraseQuery;
        }

        @Override
        public Query handle(Query query, QueryTransformer queryTransformer) {
            PhraseQuery phraseQuery = (PhraseQuery) query;
            SpanNearQuery.Builder builder = new SpanNearQuery.Builder(phraseQuery.getTerms()[0].field(), true);
            int i = 0;
            int position = -1;
            for (Term term : phraseQuery.getTerms()) {
                if (i > 0) {
                    int gap = (phraseQuery.getPositions()[i] - position) - 1;
                    if (gap > 0) {
                        builder.addGap(gap);
                    }
                }
                position = phraseQuery.getPositions()[i];
                builder.addClause(new CustomSpanPayloadCheckQuery(new SpanTermQuery(term), Collections.singletonList(null)));
                i++;
            }
            return builder.setSlop(phraseQuery.getSlop()).build();
        }
    }
}