/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF 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.mitre.quaerite.connectors;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertTrue;

import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.FutureTask;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

import org.junit.jupiter.api.Disabled;
import org.junit.jupiter.api.Test;
import org.mitre.quaerite.core.FacetResult;
import org.mitre.quaerite.core.QueryStrings;
import org.mitre.quaerite.core.SearchResultSet;
import org.mitre.quaerite.core.features.MultiMatchType;
import org.mitre.quaerite.core.features.QF;
import org.mitre.quaerite.core.features.QueryOperator;
import org.mitre.quaerite.core.features.TIE;
import org.mitre.quaerite.core.features.WeightableField;
import org.mitre.quaerite.core.queries.BooleanClause;
import org.mitre.quaerite.core.queries.BooleanQuery;
import org.mitre.quaerite.core.queries.LuceneQuery;
import org.mitre.quaerite.core.queries.MatchAllDocsQuery;
import org.mitre.quaerite.core.queries.MultiMatchQuery;
import org.mitre.quaerite.core.queries.Query;
import org.mitre.quaerite.core.queries.TermQuery;

/**
 * This class needs the tmdb collection up and running.
 *
 * Grab the template from examples and create the collection:
 *
 * curl -X PUT "http://localhost:9200/tmdb"
 *
 * curl -H "Content-Type: application/json" -T tmdb_template.json http://localhost:9200/_template/tmdb
 */


@Disabled("need to have ES tmdb instance running")
public class TestESClient {


    private static Query ALL_DOCS = new MatchAllDocsQuery();
    private static String TMDB_URL = "http://localhost:9200/tmdb";

    @Test
    public void testInit() throws Exception {
        ESClient client = new ESClient("http://localhost:9200/tmdb");
        assertEquals("http://localhost:9200/", client.getESBase());
        assertEquals("http://localhost:9200/tmdb/", client.getUrl());
        assertEquals("tmdb", client.getESCollection());
    }

    @Test
    public void testCopyFields() throws Exception {
        ESClient client = new ESClient("http://localhost:9200/tmdb");
        Collection copyFields = client.getCopyFields();
        Set<String> set = new HashSet<>(copyFields);
        assertTrue(set.contains("people"));
    }


    @Test
    public void testFacets() throws Exception {
        SearchClient client = SearchClientFactory.getClient(TMDB_URL);
        QueryRequest queryRequest = new QueryRequest(ALL_DOCS);
        queryRequest.setFacetField("genres.facet");
        queryRequest.setFacetLimit(10000);
        FacetResult result = client.facet(queryRequest);
        Map<String, Long> counts = result.getFacetCounts();
        assertEquals(21, counts.size());
        assertEquals(13582, counts.get("drama"));
        assertEquals(243, counts.get("tv movie"));


        queryRequest = new QueryRequest(ALL_DOCS);
        queryRequest.setFacetField("production_companies.facet");
        queryRequest.setFacetLimit(100000);
        result = client.facet(queryRequest);
        counts = result.getFacetCounts();
        //TODO: figure out why this # is different from Solr's
        assertEquals(15905, counts.size());
        assertEquals(1, counts.get("haile gerima"));
        assertEquals(90, counts.get("dreamworks skg"));

        queryRequest = new QueryRequest(new TermQuery("title", "red"));
        queryRequest.setFacetLimit(20000);
        queryRequest.setFacetField("production_companies.facet");
        result = client.facet(queryRequest);
        counts = result.getFacetCounts();
        assertEquals(15905, counts.size());
        assertEquals(2, counts.get("summit entertainment"));
        assertEquals(5, counts.get("paramount pictures"));
    }

    @Test
    public void testQuery() throws Exception {
        SearchClient client = SearchClientFactory.getClient(TMDB_URL);
        QueryRequest queryRequest = new QueryRequest(
                new LuceneQuery("title", "psycho"),
                null, client.getDefaultIdField());
        SearchResultSet result = client.search(queryRequest);
        Set<String> hits = new HashSet<>();
        for (int i = 0; i < result.size(); i++) {
            hits.add(result.get(i));
        }
        assertEquals(8, hits.size());
        assertTrue(hits.contains("539"));
        assertTrue(hits.contains("35683"));

        MultiMatchQuery query = new MultiMatchQuery("psycho");
        query.getQF().add(new WeightableField("title"));
        query.setMultiMatchType(new MultiMatchType("best_fields"));
        queryRequest = new QueryRequest(query,
                null, client.getDefaultIdField());
        result = client.search(queryRequest);
        hits = new HashSet<>();
        for (int i = 0; i < result.size(); i++) {
            hits.add(result.get(i));
        }
        assertEquals(8, hits.size());
        assertTrue(hits.contains("539"));
        assertTrue(hits.contains("35683"));

    }

    @Test
    public void testFilterQuery() throws Exception {
        SearchClient client = SearchClientFactory.getClient(TMDB_URL);
        LuceneQuery q = new LuceneQuery("title", "psycho");
        BooleanQuery bq = new BooleanQuery();
        bq.addClause(new BooleanClause(BooleanClause.OCCUR.MUST_NOT,
                new TermQuery("_id", "539")));
        bq.addClause(new BooleanClause(BooleanClause.OCCUR.SHOULD,
                q));
        QueryRequest queryRequest = new QueryRequest(bq,
                null, client.getDefaultIdField());

//        queryRequest.addFilterQueries(bq);
        queryRequest.setNumResults(100);
        SearchResultSet result = client.search(queryRequest);
        Set<String> hits = new HashSet<>();
        for (int i = 0; i < result.size(); i++) {
            hits.add(result.get(i));
        }
        assertEquals(7, hits.size());
        assertTrue(hits.contains("35683"));

    }

    @Test
    public void testBoolean() throws Exception {
        SearchClient client = SearchClientFactory.getClient(TMDB_URL);
        MultiMatchQuery q1 = new MultiMatchQuery("brown fox");
        q1.getQF().add(new WeightableField("title"));
        q1.getQF().add(new WeightableField("overview"));
        q1.setMultiMatchType(new MultiMatchType("best_fields"));
        q1.setTie(new TIE(0.3f));

        MultiMatchQuery q2 = new MultiMatchQuery("elephant");
        q2.getQF().add(new WeightableField("title"));
        q2.getQF().add(new WeightableField("overview"));
        q2.setMultiMatchType(new MultiMatchType("best_fields"));
        q2.setTie(new TIE(0.3f));

        BooleanQuery bq = new BooleanQuery();
        bq.addClause(new BooleanClause(BooleanClause.OCCUR.SHOULD, q1));
        bq.addClause(new BooleanClause(BooleanClause.OCCUR.SHOULD, q2));
        QueryRequest queryRequest = new QueryRequest(bq,
                null, client.getDefaultIdField());
        queryRequest.addFieldsToRetrieve("_id");
        queryRequest.setNumResults(1000);
        SearchResultSet result = client.search(queryRequest);
        Set<String> hits = new HashSet<>();
        for (int i = 0; i < result.size(); i++) {
            hits.add(result.get(i));
        }
        assertEquals(176, hits.size());
        assertTrue(hits.contains("81579"));
        assertTrue(hits.contains("42254"));

    }
    @Test
    public void testGetDocs() throws Exception {
        Set<String> ids = new HashSet<>();
        ids.addAll(Arrays.asList("539 11252 1359 10576 12662".split(" ")));
        SearchClient searchClient = SearchClientFactory.getClient(TMDB_URL);
        List<StoredDocument> docs = searchClient.getDocs("_id",
                ids, Collections.EMPTY_SET, Collections.EMPTY_SET);
        assertEquals(5, docs.size());
        StoredDocument doc1359 = null;
        for (int i = 0; i < docs.size(); i++) {
            StoredDocument sd = docs.get(i);
            if (sd.getFields().get("id").equals("1359")) {
                doc1359 = sd;
                break;
            }
        }
        String overview = (String)doc1359.getFields().get("overview");
        assertTrue(overview.startsWith("A wealthy New"));
        List<String> cast = (List)doc1359.getFields().get("cast");
        assertEquals(19, cast.size());
        assertEquals("Christian Bale", cast.get(0));
        assertEquals("Willem Dafoe", cast.get(1));
    }

    @Test
    public void testQueryOperator() throws Exception {
        QueryStrings qStrings = new QueryStrings();
        qStrings.setQuery("black mirror white christmas");
        MultiMatchQuery q = new MultiMatchQuery();
        QF qf = new QF();
        qf.add(new WeightableField("title"));
        q.setQF(qf);

        q.setMultiMatchType(new MultiMatchType("best_fields"));

        q.setQueryOperator(new QueryOperator(QueryOperator.OPERATOR.AND));
        q.setQueryStrings(qStrings);

        SearchClient searchClient = SearchClientFactory.getClient(TMDB_URL);
        SearchResultSet rs = searchClient.search(new QueryRequest(q));
        assertEquals(1, rs.getTotalHits());
        assertTrue(rs.getIds().contains("374430"));

        q.setQueryOperator(new QueryOperator(QueryOperator.OPERATOR.OR));
        rs = searchClient.search(new QueryRequest(q));
        assertEquals(369, rs.getTotalHits());
        assertTrue(rs.getIds().contains("374430"));

        q.setQueryOperator(new QueryOperator(QueryOperator.OPERATOR.OR, 0.25f));
        rs = searchClient.search(new QueryRequest(q));
        assertEquals(369, rs.getTotalHits());
        assertTrue(rs.getIds().contains("374430"));

        q.setQueryOperator(new QueryOperator(QueryOperator.OPERATOR.OR, -0.75f));
        rs = searchClient.search(new QueryRequest(q));
        assertEquals(369, rs.getTotalHits());
        assertTrue(rs.getIds().contains("374430"));

        q.setQueryOperator(new QueryOperator(QueryOperator.OPERATOR.OR, 0.5f));
        rs = searchClient.search(new QueryRequest(q));
        assertEquals(14, rs.getTotalHits());
        assertTrue(rs.getIds().contains("374430"));

        q.setQueryOperator(new QueryOperator(QueryOperator.OPERATOR.OR, 2));
        rs = searchClient.search(new QueryRequest(q));
        assertEquals(14, rs.getTotalHits());
        assertTrue(rs.getIds().contains("374430"));

        q.setQueryOperator(new QueryOperator(QueryOperator.OPERATOR.OR, -1));
        rs = searchClient.search(new QueryRequest(q));
        assertEquals(1, rs.getTotalHits());
        assertTrue(rs.getIds().contains("374430"));
    }

    @Test
    public void testGetDocsWhiteList() throws Exception {
        Set<String> ids = new HashSet<>();
        ids.addAll(Arrays.asList("539 11252 1359 10576 12662".split(" ")));
        SearchClient searchClient = SearchClientFactory.getClient(TMDB_URL);
        Set<String> whiteListFields = new HashSet<>();
        whiteListFields.add("original_title");
        List<StoredDocument> docs = searchClient.getDocs("_id", ids, whiteListFields, Collections.EMPTY_SET);
        assertEquals(5, docs.size());
        StoredDocument doc1359 = null;
        for (int i = 0; i < docs.size(); i++) {
            StoredDocument sd = docs.get(i);
            assertEquals(2, sd.getFields().size());
            if (sd.getFields().get("id").equals("1359")) {
                doc1359 = sd;
            }
        }
        String title = (String)doc1359.getFields().get("original_title");
        assertEquals("American Psycho", title);
    }

    @Test
    public void testIDGrabbing() throws Exception {
        SearchClient searchClient = SearchClientFactory.getClient(TMDB_URL);
        final ArrayBlockingQueue<Set<String>> ids = new ArrayBlockingQueue<>(10);
        IdGrabber grabber = searchClient.getIdGrabber(ids, 1000, 1, Collections.EMPTY_SET);
        Thread producer = new Thread(new FutureTask(grabber));
        producer.start();
        final AtomicInteger idCounter = new AtomicInteger(0);
        Thread consumer = new Thread() {
            @Override
            public void run() {
                while (true) {
                    Set<String> set = null;
                    try {
                        set = ids.poll(1, TimeUnit.SECONDS);
                    } catch (InterruptedException e) {
                        //swallow
                    }
                    if (set != null) {
                        idCounter.addAndGet(set.size());
                    }
                    if (set != null && set.size() == 0) {
                        break;
                    }
                }
            }
        };
        consumer.start();
        producer.join();
        consumer.join();

        assertEquals(27846, idCounter.get());
    }

    @Disabled
    public void testDeleteAll() throws Exception {
        SearchClient searchClient = SearchClientFactory.getClient(TMDB_URL);
        searchClient.deleteAll();
    }

    @Disabled("for development")
    @Test
    public void testRaw() throws Exception {
        String json =
                "{\n" +
                        "\n" +
                        "  \"query\": {\n" +
                        "  \"bool\":{\n" +
                        "    \"should\":[\n" +
                        "    {\"multi_match\" : {\n" +
                        "      \"query\":      \"brown fox\",\n" +
                        "      \"type\":       \"best_fields\",\n" +
                        "      \"fields\":     [ \"title\", \"overview\" ],\n" +
                        "      \"tie_breaker\": 0.3\n" +
                        "    }},\n" +
                        "\t{\"multi_match\" : {\n" +
                        "      \"query\":      \"elephant\",\n" +
                        "      \"type\":       \"best_fields\",\n" +
                        "      \"fields\":     [ \"title\", \"overview\" ],\n" +
                        "      \"tie_breaker\": 0.3\n" +
                        "    }}\n" +
                        "\t]\n" +
                        "  }\n" +
                        "  }\n" +
                        "}";
        SearchClient searchClient = SearchClientFactory.getClient(TMDB_URL);
        String url = TMDB_URL + "/_search";
        JsonResponse r = searchClient.postJson(url, json);
        System.out.println(r);
    }
}