import org.elasticsearch.cluster.ClusterModule;
import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
import org.elasticsearch.common.xcontent.NamedXContentRegistry;
import org.elasticsearch.common.xcontent.ParseFieldRegistry;
import org.elasticsearch.common.xcontent.XContentParser;
import org.elasticsearch.common.xcontent.json.JsonXContent;
import org.elasticsearch.index.query.AbstractQueryBuilder;
import org.elasticsearch.index.query.BoolQueryBuilder;
import org.elasticsearch.index.query.BoostingQueryBuilder;
import org.elasticsearch.index.query.CommonTermsQueryBuilder;
import org.elasticsearch.index.query.ConstantScoreQueryBuilder;
import org.elasticsearch.index.query.DisMaxQueryBuilder;
import org.elasticsearch.index.query.ExistsQueryBuilder;
import org.elasticsearch.index.query.FieldMaskingSpanQueryBuilder;
import org.elasticsearch.index.query.FuzzyQueryBuilder;
import org.elasticsearch.index.query.GeoBoundingBoxQueryBuilder;
import org.elasticsearch.index.query.GeoDistanceQueryBuilder;
import org.elasticsearch.index.query.GeoPolygonQueryBuilder;
import org.elasticsearch.index.query.IdsQueryBuilder;
import org.elasticsearch.index.query.MatchAllQueryBuilder;
import org.elasticsearch.index.query.MatchNoneQueryBuilder;
import org.elasticsearch.index.query.MatchPhrasePrefixQueryBuilder;
import org.elasticsearch.index.query.MatchPhraseQueryBuilder;
import org.elasticsearch.index.query.MatchQueryBuilder;
import org.elasticsearch.index.query.MoreLikeThisQueryBuilder;
import org.elasticsearch.index.query.MultiMatchQueryBuilder;
import org.elasticsearch.index.query.NestedQueryBuilder;
import org.elasticsearch.index.query.PrefixQueryBuilder;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.index.query.QueryStringQueryBuilder;
import org.elasticsearch.index.query.RangeQueryBuilder;
import org.elasticsearch.index.query.RegexpQueryBuilder;
import org.elasticsearch.index.query.ScriptQueryBuilder;
import org.elasticsearch.index.query.SimpleQueryStringBuilder;
import org.elasticsearch.index.query.SpanContainingQueryBuilder;
import org.elasticsearch.index.query.SpanFirstQueryBuilder;
import org.elasticsearch.index.query.SpanMultiTermQueryBuilder;
import org.elasticsearch.index.query.SpanNearQueryBuilder;
import org.elasticsearch.index.query.SpanNotQueryBuilder;
import org.elasticsearch.index.query.SpanOrQueryBuilder;
import org.elasticsearch.index.query.SpanTermQueryBuilder;
import org.elasticsearch.index.query.SpanWithinQueryBuilder;
import org.elasticsearch.index.query.TermQueryBuilder;
import org.elasticsearch.index.query.TermsQueryBuilder;
import org.elasticsearch.index.query.TermsSetQueryBuilder;
import org.elasticsearch.index.query.TypeQueryBuilder;
import org.elasticsearch.index.query.WildcardQueryBuilder;
import org.elasticsearch.index.query.WrapperQueryBuilder;
import org.elasticsearch.index.query.functionscore.FunctionScoreQueryBuilder;
import org.elasticsearch.plugins.SearchPlugin;
import org.elasticsearch.search.aggregations.AggregatorFactories;
import org.elasticsearch.search.aggregations.bucket.histogram.ParsedDateHistogram;
import org.elasticsearch.search.aggregations.bucket.significant.heuristics.SignificanceHeuristicParser;
import org.junit.Test;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

public class UnitTest {

    private final List<NamedWriteableRegistry.Entry> namedWriteables = new ArrayList<>();
    private final List<NamedXContentRegistry.Entry> namedXContents = new ArrayList<>();

    private final ParseFieldRegistry<SignificanceHeuristicParser> significanceHeuristicParserRegistry = new ParseFieldRegistry<>(
        "significance_heuristic");

    @Test
    public void term() throws IOException {
    String term = "  {\n" +
            "      \"poc_category\": {\n" +
            "        \"value\": \"sample\"\n" +
            "      }\n" +
            "    }";

    NamedXContentRegistry registry = new NamedXContentRegistry(ClusterModule.getNamedXWriteables());

    XContentParser parser = JsonXContent.jsonXContent.createParser(registry, term);
    XContentParser.Token tk = parser.nextToken();
    System.out.println(tk);

    TermQueryBuilder qb = TermQueryBuilder.fromXContent(parser);
    System.out.println("qb = " + qb.toString());
    }

    @Test
    public void term2() throws IOException {
    String term = " {\"term\": {\n" +
            "      \"poc_category\": {\n" +
            "        \"value\": \"sample\"\n" +
            "      }\n" +
            "    }}";

    registerQuery(new SearchPlugin.QuerySpec<>(TermQueryBuilder.NAME, TermQueryBuilder::new, TermQueryBuilder::fromXContent));
    NamedXContentRegistry registry = new NamedXContentRegistry(namedXContents);

    XContentParser parser = JsonXContent.jsonXContent.createParser(registry, term);
    QueryBuilder qb = AbstractQueryBuilder.parseInnerQueryBuilder(parser);
    System.out.println("qb = " + qb.toString());
    }

    @Test
    public void dateHistogram() throws IOException {
    String res = " {\n" +
            "          \"3\": {\n" +
            "            \"doc_count_error_upper_bound\": 0,\n" +
            "            \"sum_other_doc_count\": 0,\n" +
            "            \"buckets\": [\n" +
            "              {\n" +
            "                \"1\": {\n" +
            "                  \"value\": 30422\n" +
            "                },\n" +
            "                \"key\": \"PLY\",\n" +
            "                \"doc_count\": 19056\n" +
            "              },\n" +
            "              {\n" +
            "                \"1\": {\n" +
            "                  \"value\": 6\n" +
            "                },\n" +
            "                \"key\": \"MCH\",\n" +
            "                \"doc_count\": 4\n" +
            "              },\n" +
            "              {\n" +
            "                \"1\": {\n" +
            "                  \"value\": 6\n" +
            "                },\n" +
            "                \"key\": \"NBE\",\n" +
            "                \"doc_count\": 6\n" +
            "              }\n" +
            "            ]\n" +
            "          },\n" +
            "          \"key_as_string\": \"2017-12-16T00:00:00.000+09:00\",\n" +
            "          \"key\": 1513350000000,\n" +
            "          \"doc_count\": 19066\n" +
            "        }";
    System.out.println(res);


    NamedXContentRegistry registry = new NamedXContentRegistry(namedXContents);
    XContentParser parser = JsonXContent.jsonXContent.createParser(registry, res);
    ParsedDateHistogram pdh = ParsedDateHistogram.fromXContent(parser, "3");
    System.out.println("pdh = " + pdh.toString());
    }

    @Test
    public void rangeQuery() throws IOException {
    String range = "{\n" +
            "          \"range\": {\n" +
            "            \"ts\": {\n" +
            "              \"gte\": 1513263600000,\n" +
            "              \"lte\": 1513758227507,\n" +
            "              \"format\": \"epoch_millis\"\n" +
            "            }\n" +
            "          }\n" +
            "        }";

    registerQuery(new SearchPlugin.QuerySpec<>(RangeQueryBuilder.NAME, RangeQueryBuilder::new, RangeQueryBuilder::fromXContent));
    NamedXContentRegistry registry = new NamedXContentRegistry(namedXContents);

    XContentParser parser = JsonXContent.jsonXContent.createParser(registry, range);
    QueryBuilder qb = AbstractQueryBuilder.parseInnerQueryBuilder(parser);
    System.out.println("qb = " + qb.toString());
    }

    @Test
    public void aggs() throws IOException {
    String aggs = " {\n" +
            "    \"2\": {\n" +
            "      \"date_histogram\": {\n" +
            "        \"field\": \"ts\",\n" +
            "        \"interval\": \"1d\",\n" +
            "        \"time_zone\": \"Asia/Tokyo\",\n" +
            "        \"min_doc_count\": 1\n" +
            "      },\n" +
            "      \"aggs\": {\n" +
            "        \"3\": {\n" +
            "          \"terms\": {\n" +
            "            \"field\": \"log_type\",\n" +
            "            \"size\": 10,\n" +
            "            \"order\": {\n" +
            "              \"1\": \"desc\"\n" +
            "            }\n" +
            "          },\n" +
            "          \"aggs\": {\n" +
            "            \"1\": {\n" +
            "              \"sum\": {\n" +
            "                \"field\": \"datapoint\"\n" +
            "              }\n" +
            "            }\n" +
            "          }\n" +
            "        }\n" +
            "      }\n" +
            "    }\n" +
            "  }";

    System.out.println(aggs);

    NamedXContentRegistry registry = new NamedXContentRegistry(namedXContents);
    XContentParser parser = JsonXContent.jsonXContent.createParser(registry, aggs);
    AggregatorFactories.Builder ab = AggregatorFactories.parseAggregators(parser);

    System.out.println("ab = " + ab.toString());
    }

    @Test
    public void query() throws IOException {
    String q = "{\"bool\":{\"must\":[{\"query_string\":{\"query\":\"*\",\"analyze_wildcard\":true}},{\"query_string\":{\"analyze_wildcard\":true,\"query\":\"*\"}},{\"range\":{\"ts\":{\"gte\":1513263600000,\"lte\":1513752110170,\"format\":\"epoch_millis\"}}}],\"must_not\":[]}}";

    registerQuery(new SearchPlugin.QuerySpec<>(MatchQueryBuilder.NAME, MatchQueryBuilder::new, MatchQueryBuilder::fromXContent));
    registerQuery(new SearchPlugin.QuerySpec<>(MatchPhraseQueryBuilder.NAME, MatchPhraseQueryBuilder::new, MatchPhraseQueryBuilder::fromXContent));
    registerQuery(new SearchPlugin.QuerySpec<>(MatchPhrasePrefixQueryBuilder.NAME, MatchPhrasePrefixQueryBuilder::new,
            MatchPhrasePrefixQueryBuilder::fromXContent));
    registerQuery(new SearchPlugin.QuerySpec<>(MultiMatchQueryBuilder.NAME, MultiMatchQueryBuilder::new, MultiMatchQueryBuilder::fromXContent));
    registerQuery(new SearchPlugin.QuerySpec<>(NestedQueryBuilder.NAME, NestedQueryBuilder::new, NestedQueryBuilder::fromXContent));
    registerQuery(new SearchPlugin.QuerySpec<>(DisMaxQueryBuilder.NAME, DisMaxQueryBuilder::new, DisMaxQueryBuilder::fromXContent));
    registerQuery(new SearchPlugin.QuerySpec<>(IdsQueryBuilder.NAME, IdsQueryBuilder::new, IdsQueryBuilder::fromXContent));
    registerQuery(new SearchPlugin.QuerySpec<>(MatchAllQueryBuilder.NAME, MatchAllQueryBuilder::new, MatchAllQueryBuilder::fromXContent));
    registerQuery(new SearchPlugin.QuerySpec<>(QueryStringQueryBuilder.NAME, QueryStringQueryBuilder::new, QueryStringQueryBuilder::fromXContent));
    registerQuery(new SearchPlugin.QuerySpec<>(BoostingQueryBuilder.NAME, BoostingQueryBuilder::new, BoostingQueryBuilder::fromXContent));
    registerQuery(new SearchPlugin.QuerySpec<>(BoolQueryBuilder.NAME, BoolQueryBuilder::new, BoolQueryBuilder::fromXContent));
    registerQuery(new SearchPlugin.QuerySpec<>(TermQueryBuilder.NAME, TermQueryBuilder::new, TermQueryBuilder::fromXContent));
    registerQuery(new SearchPlugin.QuerySpec<>(TermsQueryBuilder.NAME, TermsQueryBuilder::new, TermsQueryBuilder::fromXContent));
    registerQuery(new SearchPlugin.QuerySpec<>(FuzzyQueryBuilder.NAME, FuzzyQueryBuilder::new, FuzzyQueryBuilder::fromXContent));
    registerQuery(new SearchPlugin.QuerySpec<>(RegexpQueryBuilder.NAME, RegexpQueryBuilder::new, RegexpQueryBuilder::fromXContent));
    registerQuery(new SearchPlugin.QuerySpec<>(RangeQueryBuilder.NAME, RangeQueryBuilder::new, RangeQueryBuilder::fromXContent));
    registerQuery(new SearchPlugin.QuerySpec<>(PrefixQueryBuilder.NAME, PrefixQueryBuilder::new, PrefixQueryBuilder::fromXContent));
    registerQuery(new SearchPlugin.QuerySpec<>(WildcardQueryBuilder.NAME, WildcardQueryBuilder::new, WildcardQueryBuilder::fromXContent));
    registerQuery(
            new SearchPlugin.QuerySpec<>(ConstantScoreQueryBuilder.NAME, ConstantScoreQueryBuilder::new, ConstantScoreQueryBuilder::fromXContent));
    registerQuery(new SearchPlugin.QuerySpec<>(SpanTermQueryBuilder.NAME, SpanTermQueryBuilder::new, SpanTermQueryBuilder::fromXContent));
    registerQuery(new SearchPlugin.QuerySpec<>(SpanNotQueryBuilder.NAME, SpanNotQueryBuilder::new, SpanNotQueryBuilder::fromXContent));
    registerQuery(new SearchPlugin.QuerySpec<>(SpanWithinQueryBuilder.NAME, SpanWithinQueryBuilder::new, SpanWithinQueryBuilder::fromXContent));
    registerQuery(new SearchPlugin.QuerySpec<>(SpanContainingQueryBuilder.NAME, SpanContainingQueryBuilder::new,
            SpanContainingQueryBuilder::fromXContent));
    registerQuery(new SearchPlugin.QuerySpec<>(FieldMaskingSpanQueryBuilder.NAME, FieldMaskingSpanQueryBuilder::new,
            FieldMaskingSpanQueryBuilder::fromXContent));
    registerQuery(new SearchPlugin.QuerySpec<>(SpanFirstQueryBuilder.NAME, SpanFirstQueryBuilder::new, SpanFirstQueryBuilder::fromXContent));
    registerQuery(new SearchPlugin.QuerySpec<>(SpanNearQueryBuilder.NAME, SpanNearQueryBuilder::new, SpanNearQueryBuilder::fromXContent));
    registerQuery(new SearchPlugin.QuerySpec<>(SpanOrQueryBuilder.NAME, SpanOrQueryBuilder::new, SpanOrQueryBuilder::fromXContent));
    registerQuery(new SearchPlugin.QuerySpec<>(MoreLikeThisQueryBuilder.NAME, MoreLikeThisQueryBuilder::new,
            MoreLikeThisQueryBuilder::fromXContent));
    registerQuery(new SearchPlugin.QuerySpec<>(WrapperQueryBuilder.NAME, WrapperQueryBuilder::new, WrapperQueryBuilder::fromXContent));
    registerQuery(new SearchPlugin.QuerySpec<>(CommonTermsQueryBuilder.NAME, CommonTermsQueryBuilder::new, CommonTermsQueryBuilder::fromXContent));
    registerQuery(
            new SearchPlugin.QuerySpec<>(SpanMultiTermQueryBuilder.NAME, SpanMultiTermQueryBuilder::new, SpanMultiTermQueryBuilder::fromXContent));
    registerQuery(new SearchPlugin.QuerySpec<>(FunctionScoreQueryBuilder.NAME, FunctionScoreQueryBuilder::new,
            FunctionScoreQueryBuilder::fromXContent));
    registerQuery(
            new SearchPlugin.QuerySpec<>(SimpleQueryStringBuilder.NAME, SimpleQueryStringBuilder::new, SimpleQueryStringBuilder::fromXContent));
    registerQuery(new SearchPlugin.QuerySpec<>(TypeQueryBuilder.NAME, TypeQueryBuilder::new, TypeQueryBuilder::fromXContent));
    registerQuery(new SearchPlugin.QuerySpec<>(ScriptQueryBuilder.NAME, ScriptQueryBuilder::new, ScriptQueryBuilder::fromXContent));
    registerQuery(new SearchPlugin.QuerySpec<>(GeoDistanceQueryBuilder.NAME, GeoDistanceQueryBuilder::new, GeoDistanceQueryBuilder::fromXContent));
    registerQuery(new SearchPlugin.QuerySpec<>(GeoBoundingBoxQueryBuilder.NAME, GeoBoundingBoxQueryBuilder::new,
            GeoBoundingBoxQueryBuilder::fromXContent));
    registerQuery(new SearchPlugin.QuerySpec<>(GeoPolygonQueryBuilder.NAME, GeoPolygonQueryBuilder::new, GeoPolygonQueryBuilder::fromXContent));
    registerQuery(new SearchPlugin.QuerySpec<>(ExistsQueryBuilder.NAME, ExistsQueryBuilder::new, ExistsQueryBuilder::fromXContent));
    registerQuery(new SearchPlugin.QuerySpec<>(MatchNoneQueryBuilder.NAME, MatchNoneQueryBuilder::new, MatchNoneQueryBuilder::fromXContent));
    registerQuery(new SearchPlugin.QuerySpec<>(TermsSetQueryBuilder.NAME, TermsSetQueryBuilder::new, TermsSetQueryBuilder::fromXContent));

    NamedXContentRegistry registry = new NamedXContentRegistry(namedXContents);

    XContentParser parser = JsonXContent.jsonXContent.createParser(registry, q);
    QueryBuilder qb = AbstractQueryBuilder.parseInnerQueryBuilder(parser);
    System.out.println("qb = " + qb.toString());
    }

    private void registerQuery(SearchPlugin.QuerySpec<?> spec) {
    namedWriteables.add(new NamedWriteableRegistry.Entry(QueryBuilder.class, spec.getName().getPreferredName(), spec.getReader()));
    namedXContents.add(new NamedXContentRegistry.Entry(QueryBuilder.class, spec.getName(),
            (p, c) -> spec.getParser().fromXContent(p)));
    }
}