package org.elasticsearch.search.fetch.subphase;

import org.elasticsearch.action.index.IndexRequestBuilder;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.index.VersionType;
import org.elasticsearch.index.query.InnerHitBuilder;
import org.elasticsearch.plugins.Plugin;
import org.elasticsearch.script.MockScriptPlugin;
import org.elasticsearch.search.SearchHits;
import org.elasticsearch.search.collapse.CollapseBuilder;
import org.elasticsearch.search.sort.FieldSortBuilder;
import org.elasticsearch.search.sort.SortOrder;
import org.elasticsearch.test.ESIntegTestCase;
import org.elasticsearch.test.InternalSettingsPlugin;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.function.Function;

import static org.elasticsearch.common.xcontent.XContentFactory.jsonBuilder;
import static org.elasticsearch.index.query.QueryBuilders.matchAllQuery;
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertHitCount;
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertNoFailures;
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertSearchHit;
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.hasId;
import static org.hamcrest.Matchers.equalTo;

public class InnerHitsCollapseIT extends ESIntegTestCase {


    @Override
    protected Collection<Class<? extends Plugin>> nodePlugins() {
        return Arrays.asList(InternalSettingsPlugin.class, CustomScriptPlugin.class);
    }

    public static class CustomScriptPlugin extends MockScriptPlugin {
        @Override
        protected Map<String, Function<Map<String, Object>, Object>> pluginScripts() {
            return Collections.singletonMap("5", script -> "5");
        }
    }

    public void testSimpleCollapse() throws Exception {
        assertAcked(prepareCreate("articles").addMapping("article", jsonBuilder().startObject().startObject("article")
                .startObject("properties")
                .startObject("title")
                  .field("type", "text")
                .endObject()
                .startObject("collapseKey")
                  .field("type", "keyword") 
                  .field("doc_values", "true")
                .endObject()
                .endObject().endObject().endObject()));

        List<IndexRequestBuilder> requests = new ArrayList<>();
        requests.add(client().prepareIndex("articles", "article", "1").setVersionType(VersionType.EXTERNAL).setVersion(111).setSource(jsonBuilder().startObject()
                .field("title", "quick brown fox")
                .field("collapseKey","ck1")                
                .endObject()));
        requests.add(client().prepareIndex("articles", "article", "2").setVersionType(VersionType.EXTERNAL).setVersion(222).setSource(jsonBuilder().startObject()
                .field("title", "big gray elephant")
                .field("collapseKey","ck2")
                .endObject()));
        requests.add(client().prepareIndex("articles", "article", "3").setVersionType(VersionType.EXTERNAL).setVersion(333).setSource(jsonBuilder().startObject()
                .field("title", "big gray mice")
                .field("collapseKey","ck1")
                .endObject()));
        indexRandom(true, requests);

        List<String> storedFields = new ArrayList<>(2);
        storedFields.add("*");
        storedFields.add("_source");
        CollapseBuilder cBuild = new CollapseBuilder("collapseKey").setInnerHits(
                new InnerHitBuilder()
                .setName("collapseKey")
                .setVersion(true)
                .setTrackScores(true)
                .setStoredFieldNames(storedFields)
                .addSort(new FieldSortBuilder("_id").order(SortOrder.ASC))
                .setFrom(0).setSize(22).setIgnoreUnmapped(false)
                );
        SearchResponse response = client().prepareSearch("articles")
                .setQuery(matchAllQuery())
                .setCollapse(cBuild)
                .setVersion(true)
                .addSort("_id", SortOrder.ASC)
                .get();
        assertNoFailures(response);
        assertHitCount(response, 3);
        assertSearchHit(response, 1, hasId("1"));
        assertEquals(response.getHits().getAt(0).getVersion(), 111L);
        assertThat(response.getHits().getAt(0).getInnerHits().size(), equalTo(1));
        SearchHits innerHits = response.getHits().getAt(0).getInnerHits().get("collapseKey");
        assertThat(innerHits.getTotalHits(), equalTo(2L));
        assertThat(innerHits.getHits().length, equalTo(2));
        assertThat(innerHits.getAt(0).getId(), equalTo("1"));
        assertEquals(innerHits.getAt(0).getVersion(), 111L);
      }



}