/*
 * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
 * or more contributor license agreements. Licensed under the Elastic License
 * 2.0 and the Server Side Public License, v 1; you may not use this file except
 * in compliance with, at your election, the Elastic License 2.0 or the Server
 * Side Public License, v 1.
 */
package org.elasticsearch.action.bulk;

import org.elasticsearch.ExceptionsHelper;
import org.elasticsearch.action.ActionListener;
import org.elasticsearch.action.DocWriteRequest;
import org.elasticsearch.action.admin.indices.refresh.RefreshRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.core.Tuple;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.rest.RestStatus;
import org.elasticsearch.test.ESIntegTestCase;

import java.util.Collections;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;

import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
import static org.hamcrest.Matchers.equalTo;
import static org.hamcrest.Matchers.lessThan;
import static org.hamcrest.Matchers.lessThanOrEqualTo;

@ESIntegTestCase.ClusterScope(scope = ESIntegTestCase.Scope.SUITE, numDataNodes = 2)
public class BulkProcessor2RetryIT extends ESIntegTestCase {
    private static final String INDEX_NAME = "test";
    Map<String, Integer> requestToExecutionCountMap = new ConcurrentHashMap<>();

    @Override
    protected Settings nodeSettings(int nodeOrdinal, Settings otherSettings) {
        // Have very low pool and queue sizes to overwhelm internal pools easily
        return Settings.builder()
            .put(super.nodeSettings(nodeOrdinal, otherSettings))
            // don't mess with this one! It's quite sensitive to a low queue size
            // (see also ThreadedActionListener which is happily spawning threads even when we already got rejected)
            // .put("thread_pool.listener.queue_size", 1)
            .put("thread_pool.get.queue_size", 1)
            // default is 200
            .put("thread_pool.write.queue_size", 30)
            .build();
    }

    public void testBulkRejectionLoadWithoutBackoff() throws Throwable {
        boolean rejectedExecutionExpected = true;
        executeBulkRejectionLoad(0, rejectedExecutionExpected);
    }

    // @TestLogging(
    // value = "org.elasticsearch.action.bulk.Retry2:trace",
    // reason = "Logging information about locks useful for tracking down deadlock"
    // )
    public void testBulkRejectionLoadWithBackoff() throws Throwable {
        boolean rejectedExecutionExpected = false;
        executeBulkRejectionLoad(8, rejectedExecutionExpected);
    }

    @SuppressWarnings("unchecked")
    private void executeBulkRejectionLoad(int maxRetries, boolean rejectedExecutionExpected) throws Throwable {
        int numberOfAsyncOps = randomIntBetween(600, 700);
        final CountDownLatch latch = new CountDownLatch(numberOfAsyncOps);
        final Set<BulkResponse> successfulResponses = Collections.newSetFromMap(new ConcurrentHashMap<>());
        final Set<Tuple<BulkRequest, Throwable>> failedResponses = Collections.newSetFromMap(new ConcurrentHashMap<>());

        assertAcked(prepareCreate(INDEX_NAME));
        ensureGreen();

        BulkProcessor2 bulkProcessor = BulkProcessor2.builder(this::countAndBulk, new BulkProcessor2.Listener() {
            @Override
            public void beforeBulk(long executionId, BulkRequest request) {
                // no op
            }

            @Override
            public void afterBulk(long executionId, BulkRequest request, BulkResponse response) {
                successfulResponses.add(response);
                latch.countDown();
            }

            @Override
            public void afterBulk(long executionId, BulkRequest request, Exception failure) {
                failedResponses.add(Tuple.tuple(request, failure));
                latch.countDown();
            }
        }, client().threadPool()).setBulkActions(1).setMaxNumberOfRetries(maxRetries).build();
        indexDocs(bulkProcessor, numberOfAsyncOps);
        latch.await(10, TimeUnit.SECONDS);
        bulkProcessor.awaitClose(1, TimeUnit.SECONDS);
        assertThat(successfulResponses.size() + failedResponses.size(), equalTo(numberOfAsyncOps));
        // validate all responses
        boolean rejectedAfterAllRetries = false;
        for (BulkResponse bulkResponse : successfulResponses) {
            for (BulkItemResponse bulkItemResponse : bulkResponse.getItems()) {
                if (bulkItemResponse.isFailed()) {
                    BulkItemResponse.Failure failure = bulkItemResponse.getFailure();
                    if (failure.getStatus() == RestStatus.TOO_MANY_REQUESTS) {
                        if (rejectedExecutionExpected == false) {
                            int count = requestToExecutionCountMap.get(bulkItemResponse.getId());
                            if (count < maxRetries + 1) {
                                throw new AssertionError("Got rejected although backoff policy would allow more retries");
                            }
                            rejectedAfterAllRetries = true;
                        }
                    } else {
                        throw new AssertionError("Unexpected failure status: " + failure.getStatus());
                    }
                }
            }
        }
        for (Tuple<BulkRequest, Throwable> failureTuple : failedResponses) {
            if (ExceptionsHelper.status(failureTuple.v2()) == RestStatus.TOO_MANY_REQUESTS) {
                if (rejectedExecutionExpected == false) {
                    for (DocWriteRequest<?> request : failureTuple.v1().requests) {
                        int count = requestToExecutionCountMap.get(request.id());
                        if (count < maxRetries + 1) {
                            throw new AssertionError("Got rejected although backoff policy would allow more retries");
                        }
                    }
                    rejectedAfterAllRetries = true;
                }
                // ignored, we exceeded the write queue size when dispatching the initial bulk request
            } else {
                Throwable t = failureTuple.v2();
                // we're not expecting any other errors
                throw new AssertionError("Unexpected failure", t);
            }
        }

        indicesAdmin().refresh(new RefreshRequest()).get();

        SearchResponse results = client().prepareSearch(INDEX_NAME).setQuery(QueryBuilders.matchAllQuery()).setSize(0).get();
        assertThat(bulkProcessor.getTotalBytesInFlight(), equalTo(0L));
        if (rejectedExecutionExpected) {
            assertThat((int) results.getHits().getTotalHits().value, lessThanOrEqualTo(numberOfAsyncOps));
        } else if (rejectedAfterAllRetries) {
            assertThat((int) results.getHits().getTotalHits().value, lessThan(numberOfAsyncOps));
        } else {
            assertThat((int) results.getHits().getTotalHits().value, equalTo(numberOfAsyncOps));
        }
    }

    private static void indexDocs(BulkProcessor2 processor, int numDocs) {
        for (int i = 1; i <= numDocs; i++) {
            processor.add(
                client().prepareIndex()
                    .setIndex(INDEX_NAME)
                    .setId(Integer.toString(i))
                    .setSource("field", randomRealisticUnicodeOfLengthBetween(1, 30))
                    .request()
            );
        }
    }

    void countAndBulk(BulkRequest request, ActionListener<BulkResponse> listener) {
        for (DocWriteRequest<?> docWriteRequest : request.requests) {
            requestToExecutionCountMap.compute(docWriteRequest.id(), (key, value) -> value == null ? 1 : value + 1);
        }
        client().bulk(request, listener);
    }

}
