/*
 * 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.apache.beam.sdk.io.elasticsearch;

import static org.apache.beam.sdk.io.elasticsearch.ElasticsearchIO.ConnectionConfiguration;
import static org.apache.beam.sdk.io.elasticsearch.ElasticsearchIO.getBackendVersion;
import static org.apache.beam.sdk.io.elasticsearch.ElasticsearchIO.parseResponse;

import com.fasterxml.jackson.databind.JsonNode;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import org.apache.http.HttpEntity;
import org.apache.http.entity.ContentType;
import org.apache.http.nio.entity.NStringEntity;
import org.elasticsearch.client.Request;
import org.elasticsearch.client.Response;
import org.elasticsearch.client.RestClient;

/** Test utilities to use with {@link ElasticsearchIO}. */
class ElasticsearchIOTestUtils {
  static final String[] FAMOUS_SCIENTISTS = {
    "Einstein",
    "Darwin",
    "Copernicus",
    "Pasteur",
    "Curie",
    "Faraday",
    "Newton",
    "Bohr",
    "Galilei",
    "Maxwell"
  };
  static final int NUM_SCIENTISTS = FAMOUS_SCIENTISTS.length;

  /** Enumeration that specifies whether to insert malformed documents. */
  public enum InjectionMode {
    INJECT_SOME_INVALID_DOCS,
    DO_NOT_INJECT_INVALID_DOCS
  }

  /** Deletes the given index synchronously. */
  static void deleteIndex(ConnectionConfiguration connectionConfiguration, RestClient restClient)
      throws IOException {
    deleteIndex(restClient, connectionConfiguration.getIndex());
  }

  private static void closeIndex(RestClient restClient, String index) throws IOException {
    restClient.performRequest("POST", String.format("/%s/_close", index));
  }

  private static void deleteIndex(RestClient restClient, String index) throws IOException {
    try {
      closeIndex(restClient, index);
      restClient.performRequest(
          "DELETE", String.format("/%s", index), Collections.singletonMap("refresh", "wait_for"));
    } catch (IOException e) {
      // it is fine to ignore this expression as deleteIndex occurs in @before,
      // so when the first tests is run, the index does not exist yet
      if (!e.getMessage().contains("index_not_found_exception")) {
        throw e;
      }
    }
  }

  /**
   * Synchronously deletes the target if it exists and then (re)creates it as a copy of source
   * synchronously.
   */
  static void copyIndex(RestClient restClient, String source, String target) throws IOException {
    deleteIndex(restClient, target);
    HttpEntity entity =
        new NStringEntity(
            String.format(
                "{\"source\" : { \"index\" : \"%s\" }, \"dest\" : { \"index\" : \"%s\" } }",
                source, target),
            ContentType.APPLICATION_JSON);
    restClient.performRequest(
        "POST", "/_reindex", Collections.singletonMap("refresh", "wait_for"), entity);
  }

  /** Inserts the given number of test documents into Elasticsearch. */
  static void insertTestDocuments(
      ConnectionConfiguration connectionConfiguration, long numDocs, RestClient restClient)
      throws IOException {
    List<String> data =
        ElasticsearchIOTestUtils.createDocuments(
            numDocs, ElasticsearchIOTestUtils.InjectionMode.DO_NOT_INJECT_INVALID_DOCS);
    StringBuilder bulkRequest = new StringBuilder();
    int i = 0;
    for (String document : data) {
      bulkRequest.append(
          String.format(
              "{ \"index\" : { \"_index\" : \"%s\", \"_type\" : \"%s\", \"_id\" : \"%s\" } }%n%s%n",
              connectionConfiguration.getIndex(),
              connectionConfiguration.getType(),
              i++,
              document));
    }
    String endPoint =
        String.format(
            "/%s/%s/_bulk", connectionConfiguration.getIndex(), connectionConfiguration.getType());
    HttpEntity requestBody =
        new NStringEntity(bulkRequest.toString(), ContentType.APPLICATION_JSON);
    Response response =
        restClient.performRequest(
            "POST", endPoint, Collections.singletonMap("refresh", "wait_for"), requestBody);
    ElasticsearchIO.checkForErrors(
        response.getEntity(), ElasticsearchIO.getBackendVersion(connectionConfiguration), false);
  }
  /**
   * Forces a refresh of the given index to make recently inserted documents available for search
   * using the index and type named in the connectionConfiguration.
   *
   * @param connectionConfiguration providing the index and type
   * @param restClient To use for issuing queries
   * @return The number of docs in the index
   * @throws IOException On error communicating with Elasticsearch
   */
  static long refreshIndexAndGetCurrentNumDocs(
      ConnectionConfiguration connectionConfiguration, RestClient restClient) throws IOException {
    return refreshIndexAndGetCurrentNumDocs(
        restClient, connectionConfiguration.getIndex(), connectionConfiguration.getType());
  }

  /**
   * Forces a refresh of the given index to make recently inserted documents available for search.
   *
   * @param restClient To use for issuing queries
   * @param index The Elasticsearch index
   * @param type The Elasticsearch type
   * @return The number of docs in the index
   * @throws IOException On error communicating with Elasticsearch
   */
  static long refreshIndexAndGetCurrentNumDocs(RestClient restClient, String index, String type)
      throws IOException {
    long result = 0;
    try {
      String endPoint = String.format("/%s/_refresh", index);
      restClient.performRequest("POST", endPoint);

      endPoint = String.format("/%s/%s/_search", index, type);
      Response response = restClient.performRequest("GET", endPoint);
      JsonNode searchResult = ElasticsearchIO.parseResponse(response.getEntity());
      result = searchResult.path("hits").path("total").asLong();
    } catch (IOException e) {
      // it is fine to ignore bellow exceptions because in testWriteWithBatchSize* sometimes,
      // we call upgrade before any doc have been written
      // (when there are fewer docs processed than batchSize).
      // In that cases index/type has not been created (created upon first doc insertion)
      if (!e.getMessage().contains("index_not_found_exception")) {
        throw e;
      }
    }
    return result;
  }

  /**
   * Generates a list of test documents for insertion.
   *
   * @param numDocs Number of docs to generate
   * @param injectionMode {@link InjectionMode} that specifies whether to insert malformed documents
   * @return the list of json String representing the documents
   */
  static List<String> createDocuments(long numDocs, InjectionMode injectionMode) {

    ArrayList<String> data = new ArrayList<>();
    for (int i = 0; i < numDocs; i++) {
      int index = i % FAMOUS_SCIENTISTS.length;
      // insert 2 malformed documents
      if (InjectionMode.INJECT_SOME_INVALID_DOCS.equals(injectionMode) && (i == 6 || i == 7)) {
        data.add(String.format("{\"scientist\";\"%s\", \"id\":%s}", FAMOUS_SCIENTISTS[index], i));
      } else {
        data.add(String.format("{\"scientist\":\"%s\", \"id\":%s}", FAMOUS_SCIENTISTS[index], i));
      }
    }
    return data;
  }

  /**
   * Executes a query for the named scientist and returns the count from the result.
   *
   * @param connectionConfiguration Specifies the index and type
   * @param restClient To use to execute the call
   * @param scientistName The scientist to query for
   * @return The cound of documents found
   * @throws IOException On error talking to Elasticsearch
   */
  static int countByScientistName(
      ConnectionConfiguration connectionConfiguration, RestClient restClient, String scientistName)
      throws IOException {
    return countByMatch(connectionConfiguration, restClient, "scientist", scientistName);
  }

  /**
   * Executes a match query for given field/value and returns the count of results.
   *
   * @param connectionConfiguration Specifies the index and type
   * @param restClient To use to execute the call
   * @param field The field to query
   * @param value The value to match
   * @return The count of documents in the search result
   * @throws IOException On error communicating with Elasticsearch
   */
  static int countByMatch(
      ConnectionConfiguration connectionConfiguration,
      RestClient restClient,
      String field,
      String value)
      throws IOException {
    String requestBody =
        "{\n"
            + "  \"query\" : {\"match\": {\n"
            + "    \""
            + field
            + "\": \""
            + value
            + "\"\n"
            + "  }}\n"
            + "}\n";
    String endPoint =
        String.format(
            "/%s/%s/_search",
            connectionConfiguration.getIndex(), connectionConfiguration.getType());
    HttpEntity httpEntity = new NStringEntity(requestBody, ContentType.APPLICATION_JSON);
    Response response =
        restClient.performRequest("GET", endPoint, Collections.emptyMap(), httpEntity);
    JsonNode searchResult = parseResponse(response.getEntity());
    return searchResult.path("hits").path("total").asInt();
  }

  public static void setIndexMapping(
      ConnectionConfiguration connectionConfiguration, RestClient restClient) throws IOException {
    String endpoint = String.format("/%s", connectionConfiguration.getIndex());
    String requestString =
        String.format(
            "{\"mappings\":{\"%s\":{\"properties\":{\"age\":{\"type\":\"long\"},"
                + " \"scientist\":{\"type\":\"%s\"}, \"id\":{\"type\":\"long\"}}}}}",
            connectionConfiguration.getType(),
            getBackendVersion(connectionConfiguration) == 2 ? "string" : "text");
    HttpEntity requestBody = new NStringEntity(requestString, ContentType.APPLICATION_JSON);
    Request request = new Request("PUT", endpoint);
    request.setEntity(requestBody);
    restClient.performRequest(request);
  }
}