/*
 * Licensed to Elasticsearch under one or more contributor
 * license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright
 * ownership. Elasticsearch licenses this file to you under
 * the Apache License, Version 2.0 (the "License"); you may
 * not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

package org.elasticsearch.action.search;

import org.apache.lucene.search.FieldComparator;
import org.apache.lucene.search.FieldDoc;
import org.apache.lucene.search.SortField;
import org.apache.lucene.search.TopFieldDocs;
import org.apache.lucene.search.TotalHits;
import org.apache.lucene.util.BytesRef;
import org.elasticsearch.common.time.DateUtils;
import org.elasticsearch.index.mapper.DateFieldMapper;
import org.elasticsearch.search.DocValueFormat;
import org.elasticsearch.test.ESTestCase;

import java.time.ZoneId;
import java.util.Arrays;

import static org.apache.lucene.search.TotalHits.Relation.GREATER_THAN_OR_EQUAL_TO;
import static org.elasticsearch.index.mapper.DateFieldMapper.DEFAULT_DATE_TIME_FORMATTER;
import static org.hamcrest.Matchers.equalTo;
import static org.hamcrest.Matchers.instanceOf;

public class BottomSortValuesCollectorTests extends ESTestCase {
    public void testWithStrings() {
        for (boolean reverse : new boolean[] { true, false }) {
            SortField[] sortFields = new SortField[] { new SortField("foo", SortField.Type.STRING_VAL, reverse) };
            DocValueFormat[] sortFormats = new DocValueFormat[] { DocValueFormat.RAW };
            BottomSortValuesCollector collector = new BottomSortValuesCollector(3, sortFields);
            collector.consumeTopDocs(createTopDocs(sortFields[0], 100,
                newBytesArray("foo", "goo", "hoo")), sortFormats);
            collector.consumeTopDocs(createTopDocs(sortFields[0], 100,
                newBytesArray("bar", "car", "zar")), sortFormats);
            collector.consumeTopDocs(createTopDocs(sortFields[0], 50,
                newBytesArray()), sortFormats);
            collector.consumeTopDocs(createTopDocs(sortFields[0], 50,
                newBytesArray("tar", "zar", "zzz")), sortFormats);
            collector.consumeTopDocs(createTopDocs(sortFields[0], 50,
                newBytesArray(null, null, "zzz")), sortFormats);
            assertThat(collector.getTotalHits(), equalTo(350L));
            assertNotNull(collector.getBottomSortValues());
            assertThat(collector.getBottomSortValues().getSortValueFormats().length, equalTo(1));
            assertThat(collector.getBottomSortValues().getSortValueFormats()[0], instanceOf(DocValueFormat.RAW.getClass()));
            assertThat(collector.getBottomSortValues().getRawSortValues().length, equalTo(1));
            assertThat(collector.getBottomSortValues().getFormattedSortValues().length, equalTo(1));
            if (reverse) {
                assertThat(collector.getBottomSortValues().getRawSortValues()[0], equalTo(new BytesRef("tar")));
                assertThat(collector.getBottomSortValues().getFormattedSortValues()[0], equalTo("tar"));
            } else {
                assertThat(collector.getBottomSortValues().getRawSortValues()[0], equalTo(new BytesRef("hoo")));
                assertThat(collector.getBottomSortValues().getFormattedSortValues()[0], equalTo("hoo"));
            }
        }
    }

    public void testWithLongs() {
        for (boolean reverse : new boolean[] { true, false }) {
            SortField[] sortFields = new SortField[] { new SortField("foo", SortField.Type.LONG, reverse) };
            DocValueFormat[] sortFormats = new DocValueFormat[]{ DocValueFormat.RAW };
            BottomSortValuesCollector collector = new BottomSortValuesCollector(3, sortFields);
            collector.consumeTopDocs(createTopDocs(sortFields[0], 100,
                newLongArray(5L, 10L, 15L)), sortFormats);
            collector.consumeTopDocs(createTopDocs(sortFields[0], 100,
                newLongArray(25L, 350L, 3500L)), sortFormats);
            collector.consumeTopDocs(createTopDocs(sortFields[0], 50,
                newLongArray(1L, 2L, 3L)), sortFormats);
            collector.consumeTopDocs(createTopDocs(sortFields[0], 50,
                newLongArray()), sortFormats);
            // ignore bottom if we have less top docs than the requested size
            collector.consumeTopDocs(createTopDocs(sortFields[0], 1,
                newLongArray(-100L)), sortFormats);
            assertNotNull(collector.getBottomSortValues());
            assertThat(collector.getTotalHits(), equalTo(301L));
            assertThat(collector.getBottomSortValues().getSortValueFormats().length, equalTo(1));
            assertThat(collector.getBottomSortValues().getSortValueFormats()[0], instanceOf(DocValueFormat.RAW.getClass()));
            assertThat(collector.getBottomSortValues().getRawSortValues().length, equalTo(1));
            assertThat(collector.getBottomSortValues().getFormattedSortValues().length, equalTo(1));
            if (reverse) {
                assertThat(collector.getBottomSortValues().getRawSortValues()[0], equalTo(25L));
                assertThat(collector.getBottomSortValues().getFormattedSortValues()[0], equalTo(25L));
            } else {
                assertThat(collector.getBottomSortValues().getRawSortValues()[0], equalTo(3L));
                assertThat(collector.getBottomSortValues().getFormattedSortValues()[0], equalTo(3L));
            }
        }
    }

    public void testWithDoubles() {
        for (boolean reverse : new boolean[] { true, false }) {
            SortField[] sortFields = new SortField[]{ new SortField("foo", SortField.Type.LONG, reverse) };
            DocValueFormat[] sortFormats = new DocValueFormat[] { DocValueFormat.RAW };
            BottomSortValuesCollector collector = new BottomSortValuesCollector(3, sortFields);
            collector.consumeTopDocs(createTopDocs(sortFields[0], 100,
                newDoubleArray(500d, 5000d, 6755d)), sortFormats);
            collector.consumeTopDocs(createTopDocs(sortFields[0], 100,
                newDoubleArray(0.1d, 1.5d, 3.5d)), sortFormats);
            collector.consumeTopDocs(createTopDocs(sortFields[0], 50,
                newDoubleArray()), sortFormats);
            collector.consumeTopDocs(createTopDocs(sortFields[0], 50,
                newDoubleArray(100d, 101d, 102d)), sortFormats);
            // ignore bottom if we have less top docs than the requested size
            collector.consumeTopDocs(createTopDocs(sortFields[0], 2,
                newDoubleArray(0d, 1d)), sortFormats);
            assertThat(collector.getTotalHits(), equalTo(302L));
            assertNotNull(collector.getBottomSortValues());
            assertThat(collector.getBottomSortValues().getSortValueFormats().length, equalTo(1));
            assertThat(collector.getBottomSortValues().getSortValueFormats()[0], instanceOf(DocValueFormat.RAW.getClass()));
            assertThat(collector.getBottomSortValues().getRawSortValues().length, equalTo(1));
            assertThat(collector.getBottomSortValues().getFormattedSortValues().length, equalTo(1));
            if (reverse) {
                assertThat(collector.getBottomSortValues().getRawSortValues()[0], equalTo(500d));
                assertThat(collector.getBottomSortValues().getFormattedSortValues()[0], equalTo(500d));
            } else {
                assertThat(collector.getBottomSortValues().getRawSortValues()[0], equalTo(3.5d));
                assertThat(collector.getBottomSortValues().getFormattedSortValues()[0], equalTo(3.5d));
            }
        }
    }

    public void testWithDates() {
        for (boolean reverse : new boolean[] { true, false }) {
            SortField[] sortFields = new SortField[]{ new SortField("foo", SortField.Type.LONG, reverse) };
            DocValueFormat[] sortFormats = new DocValueFormat[] {
                new DocValueFormat.DateTime(DEFAULT_DATE_TIME_FORMATTER, ZoneId.of("UTC"), DateFieldMapper.Resolution.MILLISECONDS)};
            BottomSortValuesCollector collector = new BottomSortValuesCollector(3, sortFields);
            collector.consumeTopDocs(createTopDocs(sortFields[0], 100,
                newDateArray("2017-06-01T12:18:20Z", "2018-04-03T15:10:27Z", "2013-06-01T13:10:20Z")), sortFormats);
            collector.consumeTopDocs(createTopDocs(sortFields[0], 100,
                newDateArray("2018-05-21T08:10:10Z", "2015-02-08T15:12:34Z", "2015-01-01T13:10:30Z")), sortFormats);
            collector.consumeTopDocs(createTopDocs(sortFields[0], 50,
                newDateArray()), sortFormats);
            collector.consumeTopDocs(createTopDocs(sortFields[0], 50,
                newDateArray("2019-12-30T07:34:20Z", "2017-03-01T12:10:30Z", "2015-07-09T14:00:30Z")), sortFormats);
            assertThat(collector.getTotalHits(), equalTo(300L));
            assertNotNull(collector.getBottomSortValues());
            assertThat(collector.getBottomSortValues().getSortValueFormats().length, equalTo(1));
            assertThat(collector.getBottomSortValues().getSortValueFormats()[0], instanceOf(DocValueFormat.DateTime.class));
            assertThat(collector.getBottomSortValues().getRawSortValues().length, equalTo(1));
            if (reverse) {
                assertThat(collector.getBottomSortValues().getRawSortValues()[0], equalTo(1436450430000L));
                assertThat(collector.getBottomSortValues().getFormattedSortValues()[0], equalTo("2015-07-09T14:00:30.000Z"));
            } else {
                assertThat(collector.getBottomSortValues().getRawSortValues()[0], equalTo(1522768227000L));
                assertThat(collector.getBottomSortValues().getFormattedSortValues()[0], equalTo("2018-04-03T15:10:27.000Z"));
            }
        }
    }

    public void testWithDateNanos() {
        for (boolean reverse : new boolean[] { true, false }) {
            SortField[] sortFields = new SortField[]{ new SortField("foo", SortField.Type.LONG, reverse) };
            DocValueFormat[] sortFormats = new DocValueFormat[] {
                new DocValueFormat.DateTime(DEFAULT_DATE_TIME_FORMATTER, ZoneId.of("UTC"), DateFieldMapper.Resolution.NANOSECONDS)};
            BottomSortValuesCollector collector = new BottomSortValuesCollector(3, sortFields);
            collector.consumeTopDocs(createTopDocs(sortFields[0], 100,
                newDateNanoArray("2017-06-01T12:18:20Z", "2018-04-03T15:10:27Z", "2013-06-01T13:10:20Z")), sortFormats);
            collector.consumeTopDocs(createTopDocs(sortFields[0], 100,
                newDateNanoArray("2018-05-21T08:10:10Z", "2015-02-08T15:12:34Z", "2015-01-01T13:10:30Z")), sortFormats);
            collector.consumeTopDocs(createTopDocs(sortFields[0], 50,
                newDateNanoArray()), sortFormats);
            collector.consumeTopDocs(createTopDocs(sortFields[0], 50,
                newDateNanoArray("2019-12-30T07:34:20Z", "2017-03-01T12:10:30Z", "2015-07-09T14:00:30Z")), sortFormats);
            assertThat(collector.getTotalHits(), equalTo(300L));
            assertNotNull(collector.getBottomSortValues());
            assertThat(collector.getBottomSortValues().getSortValueFormats().length, equalTo(1));
            assertThat(collector.getBottomSortValues().getSortValueFormats()[0], instanceOf(DocValueFormat.DateTime.class));
            assertThat(collector.getBottomSortValues().getRawSortValues().length, equalTo(1));
            if (reverse) {
                assertThat(collector.getBottomSortValues().getRawSortValues()[0], equalTo(1436450430000000000L));
                assertThat(collector.getBottomSortValues().getFormattedSortValues()[0], equalTo("2015-07-09T14:00:30.000Z"));
            } else {
                assertThat(collector.getBottomSortValues().getRawSortValues()[0], equalTo(1522768227000000000L));
                assertThat(collector.getBottomSortValues().getFormattedSortValues()[0], equalTo("2018-04-03T15:10:27.000Z"));
            }
        }
    }

    public void testWithMixedTypes() {
        for (boolean reverse : new boolean[] { true, false }) {
            SortField[] sortFields = new SortField[] { new SortField("foo", SortField.Type.LONG, reverse) };
            SortField[] otherSortFields = new SortField[] { new SortField("foo", SortField.Type.STRING_VAL, reverse) };
            DocValueFormat[] sortFormats = new DocValueFormat[] { DocValueFormat.RAW };
            BottomSortValuesCollector collector = new BottomSortValuesCollector(3, sortFields);
            collector.consumeTopDocs(createTopDocs(sortFields[0], 100,
                newLongArray(1000L, 100L, 10L)), sortFormats);
            collector.consumeTopDocs(createTopDocs(otherSortFields[0], 50,
                newBytesArray("foo", "bar", "zoo")), sortFormats);
            assertThat(collector.getTotalHits(), equalTo(150L));
            assertNotNull(collector.getBottomSortValues());
            assertThat(collector.getBottomSortValues().getSortValueFormats().length, equalTo(1));
            assertThat(collector.getBottomSortValues().getSortValueFormats()[0], instanceOf(DocValueFormat.RAW.getClass()));
            assertThat(collector.getBottomSortValues().getRawSortValues().length, equalTo(1));
            if (reverse) {
                assertThat(collector.getBottomSortValues().getRawSortValues()[0], equalTo(10L));
                assertThat(collector.getBottomSortValues().getFormattedSortValues()[0], equalTo(10L));
            } else {
                assertThat(collector.getBottomSortValues().getRawSortValues()[0], equalTo(1000L));
                assertThat(collector.getBottomSortValues().getFormattedSortValues()[0], equalTo(1000L));
            }
        }
    }

    private Object[] newDoubleArray(Double... values) {
        return values;
    }

    private Object[] newLongArray(Long... values) {
        return values;
    }

    private Object[] newBytesArray(String... values) {
        BytesRef[] bytesRefs = new BytesRef[values.length];
        for (int i = 0; i < bytesRefs.length; i++) {
            bytesRefs[i] = values[i] == null ? null : new BytesRef(values[i]);
        }
        return bytesRefs;
    }

    private Object[] newDateArray(String... values) {
        Long[] longs = new Long[values.length];
        for (int i = 0; i < values.length; i++) {
            longs[i] = DEFAULT_DATE_TIME_FORMATTER.parseMillis(values[i]);
        }
        return longs;
    }

    private Object[] newDateNanoArray(String... values) {
        Long[] longs = new Long[values.length];
        for (int i = 0; i < values.length; i++) {
            longs[i] = DateUtils.toNanoSeconds(DEFAULT_DATE_TIME_FORMATTER.parseMillis(values[i]));
        }
        return longs;
    }

    private TopFieldDocs createTopDocs(SortField sortField, int totalHits, Object[] values) {
        FieldDoc[] fieldDocs = new FieldDoc[values.length];
        FieldComparator cmp = sortField.getComparator(1, 0);
        for (int i = 0; i < values.length; i++) {
            fieldDocs[i] = new FieldDoc(i, Float.NaN, new Object[] { values[i] });
        }
        int reverseMul = sortField.getReverse() ? -1 : 1;
        Arrays.sort(fieldDocs, (o1, o2) -> reverseMul * cmp.compareValues(o1.fields[0], o2.fields[0]));
        return new TopFieldDocs(new TotalHits(totalHits, GREATER_THAN_OR_EQUAL_TO),
            fieldDocs, new SortField[] { sortField });
    }
}
