/*
 * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
 * or more contributor license agreements. Licensed under the Elastic License;
 * you may not use this file except in compliance with the Elastic License.
 */

package org.elasticsearch.xpack.flattened.mapper;

import org.elasticsearch.action.bulk.BulkRequestBuilder;
import org.elasticsearch.action.bulk.BulkResponse;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.action.support.WriteRequest.RefreshPolicy;
import org.elasticsearch.common.document.DocumentField;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.common.xcontent.XContentFactory;
import org.elasticsearch.plugins.Plugin;
import org.elasticsearch.search.aggregations.Aggregator;
import org.elasticsearch.search.aggregations.bucket.terms.Terms;
import org.elasticsearch.search.aggregations.bucket.terms.TermsAggregationBuilder;
import org.elasticsearch.search.aggregations.bucket.terms.TermsAggregatorFactory;
import org.elasticsearch.search.aggregations.metrics.Cardinality;
import org.elasticsearch.search.sort.FieldSortBuilder;
import org.elasticsearch.search.sort.SortOrder;
import org.elasticsearch.test.ESSingleNodeTestCase;
import org.elasticsearch.xpack.core.XPackPlugin;
import org.elasticsearch.xpack.flattened.FlattenedMapperPlugin;
import org.junit.Before;

import java.io.IOException;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;

import static org.elasticsearch.index.query.QueryBuilders.existsQuery;
import static org.elasticsearch.index.query.QueryBuilders.matchQuery;
import static org.elasticsearch.index.query.QueryBuilders.multiMatchQuery;
import static org.elasticsearch.index.query.QueryBuilders.queryStringQuery;
import static org.elasticsearch.index.query.QueryBuilders.simpleQueryStringQuery;
import static org.elasticsearch.search.aggregations.AggregationBuilders.cardinality;
import static org.elasticsearch.search.aggregations.AggregationBuilders.terms;
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertHitCount;
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertNoFailures;
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertOrderedSearchHits;
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertSearchResponse;
import static org.hamcrest.CoreMatchers.startsWith;
import static org.hamcrest.Matchers.equalTo;
import static org.hamcrest.Matchers.greaterThan;
import static org.hamcrest.Matchers.notNullValue;

public class FlatObjectSearchTests extends ESSingleNodeTestCase {

    protected Collection<Class<? extends Plugin>> getPlugins() {
        return pluginList(FlattenedMapperPlugin.class, XPackPlugin.class);
    }

    @Before
    public void setUpIndex() throws IOException {
        XContentBuilder mapping = XContentFactory.jsonBuilder().startObject()
            .startObject("_doc")
                .startObject("properties")
                    .startObject("flat_object")
                        .field("type", "flattened")
                        .field("split_queries_on_whitespace", true)
                    .endObject()
                    .startObject("headers")
                        .field("type", "flattened")
                        .field("split_queries_on_whitespace", true)
                    .endObject()
                    .startObject("labels")
                        .field("type", "flattened")
                    .endObject()
                .endObject()
           .endObject()
        .endObject();
        createIndex("test", Settings.EMPTY, "_doc", mapping);
    }

    public void testMatchQuery() throws Exception {
        client().prepareIndex("test", "_doc", "1")
            .setRefreshPolicy(RefreshPolicy.IMMEDIATE)
            .setSource(XContentFactory.jsonBuilder()
                .startObject()
                    .startObject("headers")
                        .field("content-type", "application/json")
                        .field("origin", "https://www.elastic.co")
                    .endObject()
            .endObject())
            .get();

        SearchResponse searchResponse = client().prepareSearch()
            .setQuery(matchQuery("headers", "application/json"))
            .get();
        assertHitCount(searchResponse, 1L);

        // Check that queries are split on whitespace.
        searchResponse = client().prepareSearch()
            .setQuery(matchQuery("headers.content-type", "application/json text/plain"))
            .get();
        assertHitCount(searchResponse, 1L);

        searchResponse = client().prepareSearch()
            .setQuery(matchQuery("headers.origin", "application/json"))
            .get();
        assertHitCount(searchResponse, 0L);
    }

    public void testMultiMatchQuery() throws Exception {
        client().prepareIndex("test", "_doc", "1")
            .setRefreshPolicy(RefreshPolicy.IMMEDIATE)
            .setSource(XContentFactory.jsonBuilder()
                .startObject()
                    .startObject("headers")
                        .field("content-type", "application/json")
                        .field("origin", "https://www.elastic.co")
                    .endObject()
            .endObject())
            .get();

        SearchResponse searchResponse = client().prepareSearch()
            .setQuery(multiMatchQuery("application/json", "headers"))
            .get();
        assertHitCount(searchResponse, 1L);

        searchResponse = client().prepareSearch()
            .setQuery(multiMatchQuery("application/json text/plain", "headers.content-type"))
            .get();
        assertHitCount(searchResponse, 1L);

        searchResponse = client().prepareSearch()
            .setQuery(multiMatchQuery("application/json", "headers.origin"))
            .get();
        assertHitCount(searchResponse, 0L);

        searchResponse = client().prepareSearch()
            .setQuery(multiMatchQuery("application/json", "headers.origin", "headers.contentType"))
            .get();
        assertHitCount(searchResponse, 0L);
    }

    public void testQueryStringQuery() throws Exception {
        client().prepareIndex("test", "_doc", "1")
            .setRefreshPolicy(RefreshPolicy.IMMEDIATE)
            .setSource(XContentFactory.jsonBuilder()
                .startObject()
                    .startObject("flat_object")
                        .field("field1", "value")
                        .field("field2", "2.718")
                    .endObject()
                .endObject())
            .get();

        SearchResponse response = client().prepareSearch("test")
            .setQuery(queryStringQuery("flat_object.field1:value"))
            .get();
        assertSearchResponse(response);
        assertHitCount(response, 1);

        response = client().prepareSearch("test")
            .setQuery(queryStringQuery("flat_object.field1:value AND flat_object:2.718"))
            .get();
        assertSearchResponse(response);
        assertHitCount(response, 1);

        response = client().prepareSearch("test")
            .setQuery(queryStringQuery("2.718").field("flat_object.field2"))
            .get();
        assertSearchResponse(response);
        assertHitCount(response, 1);
    }

    public void testSimpleQueryStringQuery() throws Exception {
        client().prepareIndex("test", "_doc", "1")
            .setRefreshPolicy(RefreshPolicy.IMMEDIATE)
            .setSource(XContentFactory.jsonBuilder()
                .startObject()
                    .startObject("flat_object")
                        .field("field1", "value")
                        .field("field2", "2.718")
                    .endObject()
                .endObject())
            .get();

        SearchResponse response = client().prepareSearch("test")
            .setQuery(simpleQueryStringQuery("value").field("flat_object.field1"))
            .get();
        assertSearchResponse(response);
        assertHitCount(response, 1);

        response = client().prepareSearch("test")
            .setQuery(simpleQueryStringQuery("+value +2.718").field("flat_object"))
            .get();
        assertSearchResponse(response);
        assertHitCount(response, 1);

        response = client().prepareSearch("test")
            .setQuery(simpleQueryStringQuery("+value +3.141").field("flat_object"))
            .get();
        assertSearchResponse(response);
        assertHitCount(response, 0);
    }

    public void testExists() throws Exception {
        client().prepareIndex("test", "_doc", "1")
            .setRefreshPolicy(RefreshPolicy.IMMEDIATE)
            .setSource(XContentFactory.jsonBuilder()
                .startObject()
                    .startObject("headers")
                        .field("content-type", "application/json")
                    .endObject()
                .endObject())
            .get();

        SearchResponse searchResponse = client().prepareSearch()
            .setQuery(existsQuery("headers"))
            .get();
        assertHitCount(searchResponse, 1L);

        searchResponse = client().prepareSearch()
            .setQuery(existsQuery("headers.content-type"))
            .get();
        assertHitCount(searchResponse, 1L);

        searchResponse = client().prepareSearch()
            .setQuery(existsQuery("headers.nonexistent"))
            .get();
        assertHitCount(searchResponse, 0L);
    }

    public void testCardinalityAggregation() throws IOException {
        int numDocs = randomIntBetween(2, 100);
        int precisionThreshold = randomIntBetween(0, 1 << randomInt(20));

        BulkRequestBuilder bulkRequest = client().prepareBulk("test", "_doc")
            .setRefreshPolicy(RefreshPolicy.IMMEDIATE);

        // Add a random number of documents containing a flat object field, plus
        // a small number of dummy documents.
        for (int i = 0; i < numDocs; ++i) {
            bulkRequest.add(client().prepareIndex()
                .setSource(XContentFactory.jsonBuilder().startObject()
                    .startObject("flat_object")
                        .field("first", i)
                        .field("second", i / 2)
                    .endObject()
                .endObject()));
        }

        for (int i = 0; i < 10; i++) {
            bulkRequest.add(client().prepareIndex("test", "_doc")
                .setSource("other_field", "1"));
        }

        BulkResponse bulkResponse = bulkRequest.get();
        assertNoFailures(bulkResponse);

        // Test the root flat object field.
        SearchResponse response = client().prepareSearch("test")
            .addAggregation(cardinality("cardinality")
                .precisionThreshold(precisionThreshold)
                .field("flat_object"))
            .get();

        assertSearchResponse(response);
        Cardinality count = response.getAggregations().get("cardinality");
        assertCardinality(count, numDocs, precisionThreshold);

        // Test two keyed flat object fields.
        SearchResponse firstResponse = client().prepareSearch("test")
            .addAggregation(cardinality("cardinality")
                .precisionThreshold(precisionThreshold)
                .field("flat_object.first"))
            .get();
        assertSearchResponse(firstResponse);

        Cardinality firstCount = firstResponse.getAggregations().get("cardinality");
        assertCardinality(firstCount, numDocs, precisionThreshold);

        SearchResponse secondResponse = client().prepareSearch("test")
            .addAggregation(cardinality("cardinality")
                .precisionThreshold(precisionThreshold)
                .field("flat_object.second"))
            .get();
        assertSearchResponse(secondResponse);

        Cardinality secondCount = secondResponse.getAggregations().get("cardinality");
        assertCardinality(secondCount, (numDocs + 1) / 2, precisionThreshold);
    }

    private void assertCardinality(Cardinality count, long value, int precisionThreshold) {
        if (value <= precisionThreshold) {
            // linear counting should be picked, and should be accurate
            assertEquals(value, count.getValue());
        } else {
            // error is not bound, so let's just make sure it is > 0
            assertThat(count.getValue(), greaterThan(0L));
        }
    }

    public void testTermsAggregation() throws IOException {
        BulkRequestBuilder bulkRequest = client().prepareBulk("test", "_doc")
            .setRefreshPolicy(RefreshPolicy.IMMEDIATE);
        for (int i = 0; i < 5; i++) {
            bulkRequest.add(client().prepareIndex()
                .setSource(XContentFactory.jsonBuilder().startObject()
                        .startObject("labels")
                            .field("priority", "urgent")
                            .field("release", "v1.2." + i)
                        .endObject()
                    .endObject()));
        }

        BulkResponse bulkResponse = bulkRequest.get();
        assertNoFailures(bulkResponse);

        // Aggregate on the root 'labels' field.
        TermsAggregationBuilder builder = createTermsAgg("labels");
        SearchResponse response = client().prepareSearch("test")
            .addAggregation(builder)
            .get();
        assertSearchResponse(response);

        Terms terms = response.getAggregations().get("terms");
        assertThat(terms, notNullValue());
        assertThat(terms.getName(), equalTo("terms"));
        assertThat(terms.getBuckets().size(), equalTo(6));

        Terms.Bucket bucket1 = terms.getBuckets().get(0);
        assertEquals("urgent", bucket1.getKey());
        assertEquals(5, bucket1.getDocCount());

        Terms.Bucket bucket2 = terms.getBuckets().get(1);
        assertThat(bucket2.getKeyAsString(), startsWith("v1.2."));
        assertEquals(1, bucket2.getDocCount());

        // Aggregate on the 'priority' subfield.
        TermsAggregationBuilder priorityAgg = createTermsAgg("labels.priority");
        SearchResponse priorityResponse = client().prepareSearch("test")
            .addAggregation(priorityAgg)
            .get();
        assertSearchResponse(priorityResponse);

        Terms priorityTerms = priorityResponse.getAggregations().get("terms");
        assertThat(priorityTerms, notNullValue());
        assertThat(priorityTerms.getName(), equalTo("terms"));
        assertThat(priorityTerms.getBuckets().size(), equalTo(1));

        Terms.Bucket priorityBucket = priorityTerms.getBuckets().get(0);
        assertEquals("urgent", priorityBucket.getKey());
        assertEquals(5, priorityBucket.getDocCount());

        // Aggregate on the 'release' subfield.
        TermsAggregationBuilder releaseAgg = createTermsAgg("labels.release");
        SearchResponse releaseResponse = client().prepareSearch("test")
            .addAggregation(releaseAgg)
            .get();
        assertSearchResponse(releaseResponse);

        Terms releaseTerms = releaseResponse.getAggregations().get("terms");
        assertThat(releaseTerms, notNullValue());
        assertThat(releaseTerms.getName(), equalTo("terms"));
        assertThat(releaseTerms.getBuckets().size(), equalTo(5));

        for (Terms.Bucket bucket : releaseTerms.getBuckets()) {
            assertThat(bucket.getKeyAsString(), startsWith("v1.2."));
            assertEquals(1, bucket.getDocCount());
        }

        // Aggregate on the 'priority' subfield with a min_doc_count of 0.
        TermsAggregationBuilder minDocCountAgg = createTermsAgg("labels.priority")
            .minDocCount(0);
        SearchResponse minDocCountResponse = client().prepareSearch("test")
            .addAggregation(minDocCountAgg)
            .get();
        assertSearchResponse(minDocCountResponse);

        Terms minDocCountTerms = minDocCountResponse.getAggregations().get("terms");
        assertThat(minDocCountTerms, notNullValue());
        assertThat(minDocCountTerms.getName(), equalTo("terms"));
        assertThat(minDocCountTerms.getBuckets().size(), equalTo(1));
    }

    private TermsAggregationBuilder createTermsAgg(String field) {
        TermsAggregatorFactory.ExecutionMode executionMode = randomFrom(
            TermsAggregatorFactory.ExecutionMode.values());
        Aggregator.SubAggCollectionMode collectionMode = randomFrom(
            Aggregator.SubAggCollectionMode.values());

        return terms("terms")
            .field(field)
            .collectMode(collectionMode)
            .executionHint(executionMode.toString());
    }


    public void testLoadDocValuesFields() throws Exception {
        client().prepareIndex("test", "_doc", "1")
            .setRefreshPolicy(RefreshPolicy.IMMEDIATE)
            .setSource(XContentFactory.jsonBuilder()
                .startObject()
                .startObject("flat_object")
                .field("key", "value")
                .field("other_key", "other_value")
                .endObject()
                .endObject())
            .get();

        SearchResponse response = client().prepareSearch("test")
            .addDocValueField("flat_object")
            .addDocValueField("flat_object.key")
            .get();
        assertSearchResponse(response);
        assertHitCount(response, 1);

        Map<String, DocumentField> fields = response.getHits().getAt(0).getFields();

        DocumentField field = fields.get("flat_object");
        assertEquals("flat_object", field.getName());
        assertEquals(Arrays.asList("other_value", "value"), field.getValues());

        DocumentField keyedField = fields.get("flat_object.key");
        assertEquals("flat_object.key", keyedField.getName());
        assertEquals("value", keyedField.getValue());
    }

    public void testFieldSort() throws Exception {
        client().prepareIndex("test", "_doc", "1")
            .setRefreshPolicy(RefreshPolicy.IMMEDIATE)
            .setSource(XContentFactory.jsonBuilder()
                .startObject()
                .startObject("flat_object")
                .field("key", "A")
                .field("other_key", "D")
                .endObject()
                .endObject())
            .get();

        client().prepareIndex("test", "_doc", "2")
            .setRefreshPolicy(RefreshPolicy.IMMEDIATE)
            .setSource(XContentFactory.jsonBuilder()
                .startObject()
                .startObject("flat_object")
                .field("key", "B")
                .field("other_key", "C")
                .endObject()
                .endObject())
            .get();

        client().prepareIndex("test", "_doc", "3")
            .setRefreshPolicy(RefreshPolicy.IMMEDIATE)
            .setSource(XContentFactory.jsonBuilder()
                .startObject()
                .startObject("flat_object")
                .field("other_key", "E")
                .endObject()
                .endObject())
            .get();

        SearchResponse response = client().prepareSearch("test")
            .addSort("flat_object", SortOrder.DESC)
            .get();
        assertSearchResponse(response);
        assertHitCount(response, 3);
        assertOrderedSearchHits(response, "3", "1", "2");

        response = client().prepareSearch("test")
            .addSort("flat_object.key", SortOrder.DESC)
            .get();
        assertSearchResponse(response);
        assertHitCount(response, 3);
        assertOrderedSearchHits(response, "2", "1", "3");

        response = client().prepareSearch("test")
            .addSort(new FieldSortBuilder("flat_object.key").order(SortOrder.DESC).missing("Z"))
            .get();
        assertSearchResponse(response);
        assertHitCount(response, 3);
        assertOrderedSearchHits(response, "3", "2", "1");
    }

    public void testSourceFiltering() {
        Map<String, Object> headers = new HashMap<>();
        headers.put("content-type", "application/json");
        headers.put("origin", "https://www.elastic.co");
        Map<String, Object> source = Collections.singletonMap("headers", headers);

        client().prepareIndex("test", "_doc", "1")
            .setRefreshPolicy(RefreshPolicy.IMMEDIATE)
            .setSource(source)
            .get();

        SearchResponse response = client().prepareSearch("test").setFetchSource(true).get();
        assertThat(response.getHits().getAt(0).getSourceAsMap(), equalTo(source));

        // Check 'include' filtering.
        response = client().prepareSearch("test").setFetchSource("headers", null).get();
        assertThat(response.getHits().getAt(0).getSourceAsMap(), equalTo(source));

        response = client().prepareSearch("test").setFetchSource("headers.content-type", null).get();
        Map<String, Object> filteredSource = Collections.singletonMap("headers",
            Collections.singletonMap("content-type", "application/json"));
        assertThat(response.getHits().getAt(0).getSourceAsMap(), equalTo(filteredSource));

        // Check 'exclude' filtering.
        response = client().prepareSearch("test").setFetchSource(null, "headers.content-type").get();
        filteredSource = Collections.singletonMap("headers",
            Collections.singletonMap("origin", "https://www.elastic.co"));
        assertThat(response.getHits().getAt(0).getSourceAsMap(), equalTo(filteredSource));
    }
}
