content: /*
   * Copyright (c) 2023 Airbyte, Inc., all rights reserved.
   */
  
  package io.airbyte.integrations.destination.elasticsearch;
  
  import co.elastic.clients.base.*;
  import co.elastic.clients.elasticsearch.ElasticsearchClient;
  import co.elastic.clients.elasticsearch._core.BulkRequest;
  import co.elastic.clients.elasticsearch._core.BulkResponse;
  import co.elastic.clients.elasticsearch._core.CreateResponse;
  import co.elastic.clients.elasticsearch._core.SearchResponse;
  import co.elastic.clients.elasticsearch._core.search.Hit;
  import co.elastic.clients.elasticsearch._core.search.HitsMetadata;
  import co.elastic.clients.elasticsearch.cat.indices.IndicesRecord;
  import co.elastic.clients.json.jackson.JacksonJsonpMapper;
  import com.fasterxml.jackson.core.JsonPointer;
  import com.fasterxml.jackson.databind.JsonNode;
  import com.fasterxml.jackson.databind.ObjectMapper;
  import io.airbyte.cdk.db.util.SSLCertificateUtils;
  import io.airbyte.protocol.models.v0.AirbyteRecordMessage;
  import jakarta.json.JsonValue;
  import java.io.IOException;
  import java.nio.charset.StandardCharsets;
  import java.util.*;
  import java.util.stream.Collectors;
  import org.apache.http.Header;
  import org.apache.http.HttpHost;
  import org.apache.http.message.BasicHeader;
  import org.elasticsearch.client.Node;
  import org.elasticsearch.client.RestClient;
  import org.elasticsearch.client.RestClientBuilder;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  /**
   * All communication with Elasticsearch should be done through this class.
   */
  public class ElasticsearchConnection {
  
    // this is the max number of hits we can query without paging
    private static final int MAX_HITS = 10000;
    private static Logger log = LoggerFactory.getLogger(ElasticsearchConnection.class);
  
    private final ElasticsearchClient client;
    private final RestClient restClient;
    private final HttpHost httpHost;
    private final ObjectMapper mapper = new ObjectMapper();
  
    /**
     * Creates a new ElasticsearchConnection that can be used to read/write records to indices
     *
     * @param config Configuration parameters for connecting to the Elasticsearch host
     */
    public ElasticsearchConnection(ConnectorConfiguration config) {
      log.info(String.format(
          "creating ElasticsearchConnection: %s", config.getEndpoint()));
  
      // Create the low-level client
      httpHost = HttpHost.create(config.getEndpoint());
      final RestClientBuilder builder;
      if (config.getPathPrefix() != null) {
        builder = RestClient.builder(httpHost).setPathPrefix(config.getPathPrefix());
      } else {
        builder = RestClient.builder(httpHost);
      }
  
      // Set custom user's certificate if provided
      if (config.getCaCertificate() != null && !config.getCaCertificate().isEmpty()) {
        builder.setHttpClientConfigCallback(clientBuilder -> {
          clientBuilder.setSSLContext(SSLCertificateUtils.createContextFromCaCert(config.getCaCertificate()));
          return clientBuilder;
        });
      }
  
      restClient = builder
          .setDefaultHeaders(configureHeaders(config))
          .setFailureListener(new FailureListener())
          .build();
      // Create the transport that provides JSON and http services to API clients
      Transport transport = new RestClientTransport(restClient, new JacksonJsonpMapper());
      // And create our API client
      client = new ElasticsearchClient(transport);
    }
  
    static class FailureListener extends RestClient.FailureListener {
  
      @Override
      public void onFailure(Node node) {
        log.error("RestClient failure: {}", node);
      }
  
    }
  
    /**
     * Configures the default headers for requests to the Elasticsearch server
     *
     * @param config connection information
     * @return the default headers
     */
    protected Header[] configureHeaders(ConnectorConfiguration config) {
      final var headerList = new ArrayList<Header>();
      // add Authorization header if credentials are present
      final var auth = config.getAuthenticationMethod();
      switch (auth.getMethod()) {
        case secret -> {
          var bytes = (auth.getApiKeyId() + ":" + auth.getApiKeySecret()).getBytes(StandardCharsets.UTF_8);
          var header = "ApiKey " + Base64.getEncoder().encodeToString(bytes);
          headerList.add(new BasicHeader("Authorization", header));
        }
        case basic -> {
          var basicBytes = (auth.getUsername() + ":" + auth.getPassword()).getBytes(StandardCharsets.UTF_8);
          var basicHeader = "Basic " + Base64.getEncoder().encodeToString(basicBytes);
          headerList.add(new BasicHeader("Authorization", basicHeader));
        }
      }
      return headerList.toArray(new Header[headerList.size()]);
    }
  
    /**
     * Pings the Elasticsearch server for "up" check, and configuration validation
     *
     * @return true if connection was successful
     */
    public boolean checkConnection() {
      log.info("checking elasticsearch connection");
      try {
        final var info = client.info();
        log.info("checked elasticsearch connection: {}, version: {}", info.clusterName(), info.version());
        return true;
      } catch (ApiException e) {
        log.error("failed to ping elasticsearch", unwrappedApiException("failed write operation", e));
        return false;
      } catch (Exception e) {
        log.error("unknown exception while pinging elasticsearch server", e);
        return false;
      }
    }
  
    /**
     * Writes a single record to the Elasticsearch server
     *
     * @param index The index to write the record to
     * @param id The ID to give the new document
     * @param data The body of the document (source record)
     * @return results of the create operation
     * @throws Exception if an error is encountered
     */
    public CreateResponse createDocument(String index, String id, JsonNode data) throws Exception {
      CreateResponse createResponse = client.create(builder -> builder.id(id).document(data).index(index));
      log.debug("wrote record: {}", createResponse.result());
      return createResponse;
    }
  
    /**
     * Bulk operation to append multiple documents to an Elasticsearch server
     *
     * @param index The index to add the documents to
     * @param records The collection of records to create documents from
     * @return the response of the bulk operation
     * @throws IOException if there is server connection problem, or a non-successful operation on the
     *         server
     */
    public BulkResponse indexDocuments(String index, List<AirbyteRecordMessage> records, ElasticsearchWriteConfig config) throws IOException {
      var bulkRequest = new BulkRequest.Builder<>();
      for (var doc : records) {
        log.debug("adding record to bulk create: {}", doc.getData());
        bulkRequest.addOperation(
            b -> b.index(
                c -> c.index(index).id(extractPrimaryKey(doc, config))))
            .addDocument(doc.getData()).refresh(JsonValue.TRUE);
      }
  
      try {
        return client.bulk(b -> bulkRequest);
      } catch (ApiException e) {
        throw unwrappedApiException("failed write operation", e);
      }
    }
  
    // TODO: Can we do something like this?
    private String extractPrimaryKey(AirbyteRecordMessage doc, ElasticsearchWriteConfig config) {
      if (!config.hasPrimaryKey()) {
        return UUID.randomUUID().toString();
      }
      var optFirst = config.getPrimaryKey().stream().findFirst();
      StringBuilder sb = new StringBuilder();
      if (optFirst.isPresent()) {
        log.debug("trying to extract primary key using {}", optFirst.get());
        optFirst.get().forEach(s -> sb.append(String.format("/%s", s)));
      }
      if (sb.length() > 0) {
        JsonPointer ptr = JsonPointer.valueOf(sb.toString());
        var pkNode = doc.getData().at(ptr);
        if (!pkNode.isMissingNode() && pkNode.isValueNode()) {
          return pkNode.asText();
        }
      }
      log.warn("unable to extract primary key");
      return UUID.randomUUID().toString();
    }
  
    /**
     * returns the first 10k documents of a given index
     *
     * @param index the index to search
     * @return a list of matching documents
     * @throws IOException if there is server communication error, or invalid index
     */
    public List<JsonNode> getRecords(String index) throws IOException {
      log.info("getting records for index: {}", index);
      SearchResponse<JsonNode> search = client.search(s -> s.index(index).size(MAX_HITS), JsonNode.class);
      HitsMetadata<JsonNode> hitMeta = search.hits();
      return hitMeta.hits().stream().map(Hit::source).collect(Collectors.toList());
    }
  
    /**
     * Shutdown the connection to the Elasticsearch server
     */
    public void close() throws IOException {
      this.restClient.close();
      this.client.shutdown();
    }
  
    public List<String> allIndices() {
      try {
        var response = client.cat().indices(r -> r);
        return response.valueBody().stream().map(IndicesRecord::index).collect(Collectors.toList());
      } catch (ApiException e) {
        log.error("", unwrappedApiException("failed to get indices", e));
      } catch (IOException e) {
        log.error("unknown exception while getting indices", e);
      }
      return new ArrayList<>();
    }
  
    /**
     * Creates an index on Elasticsearch if it's missing
     *
     * @param index the index name to create
     */
    public void createIndexIfMissing(String index) {
      try {
        BooleanResponse existsResponse = client.indices().exists(b -> b.index(index));
        if (existsResponse.value()) {
          log.info("index exists: {}", index);
          return;
        }
        log.info("creating index: {}, info: {}", index, client.info());
        final co.elastic.clients.elasticsearch.indices.CreateResponse createResponse = client.indices().create(b -> b.index(index));
        if (createResponse.acknowledged() && createResponse.shardsAcknowledged()) {
          log.info("created index: {}", index);
        } else {
          log.info("did not create index: {}, {}", index, mapper.writeValueAsString(createResponse));
        }
      } catch (ApiException e) {
        log.warn("", unwrappedApiException("failed to create index", e));
      } catch (IOException e) {
        log.warn("unknown exception while creating index", e);
      }
    }
  
    /**
     * Deletes an index if present, suppressing any exceptions
     *
     * @param indexName The index to delete
     */
    public void deleteIndexIfPresent(String indexName) {
      try {
        var exists = client.indices().exists(b -> b.index(indexName));
        if (exists.value()) {
          client.indices().delete(b -> b.index(indexName));
        }
      } catch (ApiException e) {
        log.warn("", unwrappedApiException("failed to delete index", e));
      } catch (IOException e) {
        log.warn("unknown exception while deleting index", e);
      }
    }
  
    /**
     * Clones a source index to a destination index. If the destination index already exists, it deletes
     * it before cloning
     *
     * @param sourceIndexName The index to clone
     * @param destinationIndexName The destination index name to clone to.
     */
    public void replaceIndex(String sourceIndexName, String destinationIndexName) {
      log.info("replacing index: {}, with index: {}", destinationIndexName, sourceIndexName);
      try {
        var sourceExists = client.indices().exists(i -> i.index(sourceIndexName));
        if (!sourceExists.value()) {
          throw new RuntimeException(
              String.format("the source index does not exist. unable to replace the destination index. source: %s, destination: %s", sourceIndexName,
                  destinationIndexName));
        }
  
        // delete the destination if it exists
        var destinationExists = client.indices().exists(i -> i.index(destinationIndexName));
        if (destinationExists.value()) {
          log.warn("deleting existing index: {}", destinationIndexName);
          deleteIndexIfPresent(destinationIndexName);
        }
  
        // make the source index read-only before cloning
        // I think theres a bug in the client
        // https://github.com/elastic/elasticsearch-java/issues/37
        // client.indices().addBlock(b -> b.index(sourceIndexName).block(IndicesBlockOptions.Write));
        // so we need to do it a different way
        client.indices().putSettings(b -> b.index(sourceIndexName).settings(s -> s.blocks(w -> w.write(true))));
  
        // clone to the destination
        client.indices().clone(c -> c.index(sourceIndexName).target(destinationIndexName));
  
        // enable writing on new index
        client.indices().putSettings(b -> b.index(destinationIndexName).settings(s -> s.blocks(w -> w.write(false))));
      } catch (ApiException e) {
        throw unwrappedApiException("failed to delete index", e);
      } catch (IOException e) {
        throw new RuntimeException("unknown exception while replacing index", e);
      }
    }
  
    /**
     * Unwraps a rest client ApiException, so we can log the details
     *
     * @param message message to add to the log entry
     * @param e source ApiException
     * @return a new RuntimeException with the ApiException as the source
     */
    private RuntimeException unwrappedApiException(String message, ApiException e) {
      log.error(message);
      if (Objects.isNull(e) || Objects.isNull(e.error())) {
        log.error("unknown ApiException");
        return new RuntimeException(e);
      }
      if (ElasticsearchError.class.isAssignableFrom(e.error().getClass())) {
        ElasticsearchError esException = ((ElasticsearchError) e.error());
        String errorMessage = String.format("ElasticsearchError: status:%s, error:%s", esException.status(), esException.error().toString());
        return new RuntimeException(errorMessage);
      }
      return new RuntimeException(e);
    }
  
  }
  