/*
 * 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.samza.system.elasticsearch;

import org.apache.samza.SamzaException;
import org.apache.samza.system.OutgoingMessageEnvelope;
import org.apache.samza.system.SystemProducer;
import org.apache.samza.system.elasticsearch.indexrequest.IndexRequestFactory;
import org.elasticsearch.action.ActionResponse;
import org.elasticsearch.action.bulk.BulkItemResponse;
import org.elasticsearch.action.bulk.BulkProcessor;
import org.elasticsearch.action.bulk.BulkRequest;
import org.elasticsearch.action.bulk.BulkResponse;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.action.index.IndexResponse;
import org.elasticsearch.client.Client;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Arrays;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;

/** A {@link SystemProducer} for Elasticsearch that builds on top of the {@link BulkProcessor}
 *
 * <p>
 * Each system that is configured in Samza has an independent {@link BulkProcessor} that flush
 * separably to Elasticsearch. Each {@link BulkProcessor} will maintain the ordering of messages
 * being sent from tasks per Samza container. If you have multiple containers writing to the same
 * message id there is no guarantee of ordering in Elasticsearch.
 * </p>
 *
 * <p>
 * This can be fully configured from the Samza job properties. The client factory and index request
 * are pluggable so the implementation of these can be changed if required.
 * </p>
 *
 * */
public class ElasticsearchSystemProducer implements SystemProducer {
  private static final Logger LOGGER = LoggerFactory.getLogger(ElasticsearchSystemProducer.class);

  private final String system;
  private final Map<String, BulkProcessor> sourceBulkProcessor;
  private final AtomicBoolean sendFailed = new AtomicBoolean(false);
  private final AtomicReference<Throwable> thrown = new AtomicReference<>();

  private final IndexRequestFactory indexRequestFactory;
  private final BulkProcessorFactory bulkProcessorFactory;
  private final ElasticsearchSystemProducerMetrics metrics;

  private Client client;

  public ElasticsearchSystemProducer(String system, BulkProcessorFactory bulkProcessorFactory,
                                     Client client, IndexRequestFactory indexRequestFactory,
                                     ElasticsearchSystemProducerMetrics metrics) {
    this.system = system;
    this.sourceBulkProcessor = new HashMap<>();
    this.bulkProcessorFactory = bulkProcessorFactory;
    this.client = client;
    this.indexRequestFactory = indexRequestFactory;
    this.metrics = metrics;
  }


  @Override
  public void start() {
    // Nothing to do.
  }

  @Override
  public void stop() {
    for (Map.Entry<String, BulkProcessor> e : sourceBulkProcessor.entrySet()) {
      flush(e.getKey());
      e.getValue().close();
    }

    client.close();
  }

  @Override
  public void register(final String source) {
    BulkProcessor.Listener listener = new BulkProcessor.Listener() {
        @Override
        public void beforeBulk(long executionId, BulkRequest request) {
          // Nothing to do.
        }

        @Override
        public void afterBulk(long executionId, BulkRequest request, BulkResponse response) {
          int respLen = response.getItems().length;

          if (response.hasFailures()) {
            sendFailed.set(true);
          } else {
            updateSuccessMetrics(response);
            LOGGER.info(String.format("Written %s messages from %s to %s.",
                    response.getItems().length, source, system));
          }
        }

        @Override
        public void afterBulk(long executionId, BulkRequest request, Throwable failure) {
          thrown.compareAndSet(null, failure);
          sendFailed.set(true);
        }
    };

    sourceBulkProcessor.put(source, bulkProcessorFactory.getBulkProcessor(client, listener));
  }

  private void updateSuccessMetrics(BulkResponse response) {
    metrics.bulkSendSuccess.inc();
    for (BulkItemResponse itemResp: response.getItems()) {
        ActionResponse resp = itemResp.getResponse();
        if (resp instanceof IndexResponse) {
          if (((IndexResponse)resp).isCreated()) {
            metrics.inserts.inc();
          }
          else {
            metrics.updates.inc();
          }
        }
      }
  }

  @Override
  public void send(String source, OutgoingMessageEnvelope envelope) {
    IndexRequest indexRequest = indexRequestFactory.getIndexRequest(envelope);
    sourceBulkProcessor.get(source).add(indexRequest);
  }

  @Override
  public void flush(String source) {
    sourceBulkProcessor.get(source).flush();

    if (sendFailed.get()) {
      String message = String.format("Unable to send message from %s to system %s.", source,
                                     system);
      LOGGER.error(message);

      Throwable cause = thrown.get();
      if (cause != null) {
        throw new SamzaException(message, cause);
      } else {
        throw new SamzaException(message);
      }
    }

    LOGGER.info(String.format("Flushed %s to %s.", source, system));
  }
}

@@ -1535,33 +1535,8 @@
 y;%0A%0A
-import java.util.Arrays;%0A
 impo
@@ -4114,61 +4114,8 @@
 ) %7B%0A
-          int respLen = response.getItems().length;%0A%0A
     

Clean up