/*
 * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
 * or more contributor license agreements. Licensed under the Elastic License;
 * you may not use this file except in compliance with the Elastic License.
 */
package org.elasticsearch.xpack.eql.async;

import org.elasticsearch.action.ActionListener;
import org.elasticsearch.action.ActionRequest;
import org.elasticsearch.action.ActionRequestValidationException;
import org.elasticsearch.action.ActionResponse;
import org.elasticsearch.cluster.service.ClusterService;
import org.elasticsearch.common.io.stream.StreamInput;
import org.elasticsearch.common.io.stream.StreamOutput;
import org.elasticsearch.common.unit.TimeValue;
import org.elasticsearch.tasks.Task;
import org.elasticsearch.tasks.TaskId;
import org.elasticsearch.test.ESSingleNodeTestCase;
import org.elasticsearch.transport.TransportService;
import org.elasticsearch.xpack.core.async.AsyncExecutionId;
import org.elasticsearch.xpack.core.async.AsyncResultsService;
import org.elasticsearch.xpack.core.async.AsyncTaskIndexService;
import org.elasticsearch.xpack.core.async.GetAsyncResultRequest;
import org.junit.After;
import org.junit.Before;

import java.io.IOException;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;

import static org.elasticsearch.xpack.eql.async.AsyncTaskManagementService.addCompletionListener;
import static org.hamcrest.Matchers.equalTo;
import static org.hamcrest.Matchers.notNullValue;
import static org.hamcrest.Matchers.nullValue;

public class AsyncTaskManagementServiceTests extends ESSingleNodeTestCase {
    private ClusterService clusterService;
    private TransportService transportService;
    private AsyncResultsService<TestTask, StoredAsyncResponse<TestResponse>> results;

    private final ExecutorService executorService = Executors.newFixedThreadPool(1);

    public static class TestRequest extends ActionRequest {
        private final String string;

        public TestRequest(String string) {
            this.string = string;
        }

        @Override
        public ActionRequestValidationException validate() {
            return null;
        }
    }

    public static class TestResponse extends ActionResponse {
        private final String string;
        private final String id;

        public TestResponse(String string, String id) {
            this.string = string;
            this.id = id;
        }

        public TestResponse(StreamInput input) throws IOException {
            this.string = input.readOptionalString();
            this.id = input.readOptionalString();
        }

        @Override
        public void writeTo(StreamOutput out) throws IOException {
            out.writeOptionalString(string);
            out.writeOptionalString(id);
        }
    }

    public static class TestTask extends StoredAsyncTask<TestResponse> {
        public volatile AtomicReference<TestResponse> finalResponse = new AtomicReference<>();

        public TestTask(long id, String type, String action, String description, TaskId parentTaskId, Map<String, String> headers,
                        Map<String, String> originHeaders, AsyncExecutionId asyncExecutionId, TimeValue keepAlive) {
            super(id, type, action, description, parentTaskId, headers, originHeaders, asyncExecutionId, keepAlive);
        }

        @Override
        public TestResponse getCurrentResult() {
            return Objects.requireNonNullElseGet(finalResponse.get(), () -> new TestResponse(null, getExecutionId().getEncoded()));
        }
    }

    public static class TestOperation implements AsyncTaskManagementService.AsyncOperation<TestRequest, TestResponse, TestTask> {

        @Override
        public TestTask createTask(TestRequest request, long id, String type, String action, TaskId parentTaskId,
                                   Map<String, String> headers, Map<String, String> originHeaders, AsyncExecutionId asyncExecutionId) {
            return new TestTask(id, type, action, request.getDescription(), parentTaskId, headers, originHeaders, asyncExecutionId,
                TimeValue.timeValueDays(5));
        }

        @Override
        public void execute(TestRequest request, TestTask task, ActionListener<TestResponse> listener) {
            if (request.string.equals("die")) {
                listener.onFailure(new IllegalArgumentException("test exception"));
            } else {
                listener.onResponse(new TestResponse("response for [" + request.string + "]", task.getExecutionId().getEncoded()));
            }
        }

        @Override
        public TestResponse initialResponse(TestTask task) {
            return new TestResponse(null, task.getExecutionId().getEncoded());
        }

        @Override
        public TestResponse readResponse(StreamInput inputStream) throws IOException {
            return new TestResponse(inputStream);
        }
    }

    public String index = "test-index";

    @Before
    public void setup() {
        clusterService = getInstanceFromNode(ClusterService.class);
        transportService = getInstanceFromNode(TransportService.class);
        AsyncTaskIndexService<StoredAsyncResponse<TestResponse>> store =
            new AsyncTaskIndexService<>(index, clusterService, transportService.getThreadPool().getThreadContext(), client(), "test",
                in -> new StoredAsyncResponse<>(TestResponse::new, in), writableRegistry());
        results = new AsyncResultsService<>(store, true, TestTask.class,
            (task, listener, timeout) -> addCompletionListener(transportService.getThreadPool(), task, listener, timeout),
            transportService.getTaskManager(), clusterService);
    }

    /**
     * Shutdown the executor so we don't leak threads into other test runs.
     */
    @After
    public void shutdownExec() {
        executorService.shutdown();
    }

    private AsyncTaskManagementService<TestRequest, TestResponse, TestTask> createManagementService(
        AsyncTaskManagementService.AsyncOperation<TestRequest, TestResponse, TestTask> operation) {
        return new AsyncTaskManagementService<>(index, client(), "test_origin", writableRegistry(),
            transportService.getTaskManager(), "test_action", operation, TestTask.class, clusterService, transportService.getThreadPool());
    }

    public void testReturnBeforeTimeout() throws Exception {
        AsyncTaskManagementService<TestRequest, TestResponse, TestTask> service = createManagementService(new TestOperation());
        boolean success = randomBoolean();
        boolean keepOnCompletion = randomBoolean();
        CountDownLatch latch = new CountDownLatch(1);
        TestRequest request = new TestRequest(success ? randomAlphaOfLength(10) : "die");
        service.asyncExecute(request, TimeValue.timeValueMinutes(1), TimeValue.timeValueMinutes(10), keepOnCompletion,
            ActionListener.wrap(r -> {
                assertThat(success, equalTo(true));
                assertThat(r.string, equalTo("response for [" + request.string + "]"));
                assertThat(r.id, notNullValue());
                latch.countDown();
            }, e -> {
                assertThat(success, equalTo(false));
                assertThat(e.getMessage(), equalTo("test exception"));
                latch.countDown();
            }));
        assertThat(latch.await(10, TimeUnit.SECONDS), equalTo(true));
    }

    public void testReturnAfterTimeout() throws Exception {
        CountDownLatch executionLatch = new CountDownLatch(1);
        AsyncTaskManagementService<TestRequest, TestResponse, TestTask> service = createManagementService(new TestOperation() {
            @Override
            public void execute(TestRequest request, TestTask task, ActionListener<TestResponse> listener) {
                executorService.submit(() -> {
                    try {
                        assertThat(executionLatch.await(10, TimeUnit.SECONDS), equalTo(true));
                    } catch (InterruptedException ex) {
                        fail("Shouldn't be here");
                    }
                    super.execute(request, task, listener);
                });
            }
        });
        boolean success = randomBoolean();
        boolean keepOnCompletion = randomBoolean();
        boolean timeoutOnFirstAttempt = randomBoolean();
        boolean waitForCompletion = randomBoolean();
        CountDownLatch latch = new CountDownLatch(1);
        TestRequest request = new TestRequest(success ? randomAlphaOfLength(10) : "die");
        AtomicReference<TestResponse> responseHolder = new AtomicReference<>();
        service.asyncExecute(request, TimeValue.timeValueMillis(1), TimeValue.timeValueMinutes(10), keepOnCompletion,
            ActionListener.wrap(r -> {
                assertThat(r.string, nullValue());
                assertThat(r.id, notNullValue());
                assertThat(responseHolder.getAndSet(r), nullValue());
                latch.countDown();
            }, e -> fail("Shouldn't be here")));
        assertThat(latch.await(20, TimeUnit.SECONDS), equalTo(true));

        if (timeoutOnFirstAttempt) {
            logger.trace("Getting an in-flight response");
            // try getting results, but fail with timeout because it is not ready yet
            StoredAsyncResponse<TestResponse> response = getResponse(responseHolder.get().id, TimeValue.timeValueMillis(2));
            assertThat(response.getException(), nullValue());
            assertThat(response.getResponse(), notNullValue());
            assertThat(response.getResponse().id, equalTo(responseHolder.get().id));
            assertThat(response.getResponse().string, nullValue());
        }

        if (waitForCompletion) {
            // now we are waiting for the task to finish
            logger.trace("Waiting for response to complete");
            AtomicReference<StoredAsyncResponse<TestResponse>> responseRef = new AtomicReference<>();
            CountDownLatch getResponseCountDown = getResponse(responseHolder.get().id, TimeValue.timeValueSeconds(5),
                ActionListener.wrap(responseRef::set, e -> fail("Shouldn't be here")));

            executionLatch.countDown();
            assertThat(getResponseCountDown.await(10, TimeUnit.SECONDS), equalTo(true));

            StoredAsyncResponse<TestResponse> response = responseRef.get();
            if (success) {
                assertThat(response.getException(), nullValue());
                assertThat(response.getResponse(), notNullValue());
                assertThat(response.getResponse().id, equalTo(responseHolder.get().id));
                assertThat(response.getResponse().string, equalTo("response for [" + request.string + "]"));
            } else {
                assertThat(response.getException(), notNullValue());
                assertThat(response.getResponse(), nullValue());
                assertThat(response.getException().getMessage(), equalTo("test exception"));
            }
        } else {
            executionLatch.countDown();
        }

        // finally wait until the task disappears and get the response from the index
        logger.trace("Wait for task to disappear ");
        assertBusy(() -> {
            Task task = transportService.getTaskManager().getTask(AsyncExecutionId.decode(responseHolder.get().id).getTaskId().getId());
            assertThat(task, nullValue());
        });

        logger.trace("Getting the the final response from the index");
        StoredAsyncResponse<TestResponse> response = getResponse(responseHolder.get().id, TimeValue.ZERO);
        if (success) {
            assertThat(response.getException(), nullValue());
            assertThat(response.getResponse(), notNullValue());
            assertThat(response.getResponse().string, equalTo("response for [" + request.string + "]"));
        } else {
            assertThat(response.getException(), notNullValue());
            assertThat(response.getResponse(), nullValue());
            assertThat(response.getException().getMessage(), equalTo("test exception"));
        }
    }

    private StoredAsyncResponse<TestResponse> getResponse(String id, TimeValue timeout) throws InterruptedException {
        AtomicReference<StoredAsyncResponse<TestResponse>> response = new AtomicReference<>();
        assertThat(
            getResponse(id, timeout, ActionListener.wrap(response::set, e -> fail("Shouldn't be here"))).await(10, TimeUnit.SECONDS),
            equalTo(true)
        );
        return response.get();
    }

    private CountDownLatch getResponse(String id,
                                       TimeValue timeout,
                                       ActionListener<StoredAsyncResponse<TestResponse>> listener) {
        CountDownLatch responseLatch = new CountDownLatch(1);
        GetAsyncResultRequest getResultsRequest = new GetAsyncResultRequest(id)
            .setWaitForCompletionTimeout(timeout);
        results.retrieveResult(getResultsRequest, ActionListener.wrap(
            r -> {
                listener.onResponse(r);
                responseLatch.countDown();
            },
            e -> {
                listener.onFailure(e);
                responseLatch.countDown();
            }
        ));
        return responseLatch;
    }

}
