/**
 * 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
 * <p>
 * http://www.apache.org/licenses/LICENSE-2.0
 * <p>
 * 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.atlas.notification;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.util.concurrent.ThreadFactoryBuilder;
import org.apache.atlas.ApplicationProperties;
import org.apache.atlas.AtlasException;
import org.apache.atlas.exception.AtlasBaseException;
import org.apache.atlas.ha.HAConfiguration;
import org.apache.atlas.listener.ActiveStateChangeHandler;
import org.apache.atlas.model.impexp.AtlasAsyncImportRequest;
import org.apache.atlas.model.impexp.AtlasAsyncImportRequest.ImportStatus;
import org.apache.atlas.repository.impexp.AsyncImportService;
import org.apache.atlas.repository.store.graph.v2.asyncimport.ImportTaskListener;
import org.apache.atlas.service.Service;
import org.apache.commons.configuration.Configuration;
import org.apache.commons.lang3.ObjectUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.context.annotation.DependsOn;
import org.springframework.core.annotation.Order;
import org.springframework.stereotype.Component;

import javax.annotation.PreDestroy;
import javax.inject.Inject;

import java.util.List;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.Semaphore;
import java.util.concurrent.TimeUnit;
import java.util.stream.Stream;

import static org.apache.atlas.AtlasConfiguration.ASYNC_IMPORT_TOPIC_PREFIX;
import static org.apache.atlas.AtlasErrorCode.IMPORT_QUEUEING_FAILED;

@Component
@Order(8)
@DependsOn(value = "notificationHookConsumer")
public class ImportTaskListenerImpl implements Service, ActiveStateChangeHandler, ImportTaskListener {
    private static final Logger            LOG                  = LoggerFactory.getLogger(ImportTaskListenerImpl.class);

    private static final String            THREADNAME_PREFIX    = ImportTaskListener.class.getSimpleName();
    private static final int               ASYNC_IMPORT_PERMITS = 1; // Only one asynchronous import task is permitted

    private volatile boolean               isActiveInstance     = true;
    private volatile ExecutorService       executorService; // Single-thread executor for sequential processing
    private final BlockingQueue<String>    requestQueue;    // Blocking queue for requests
    private final AsyncImportService       asyncImportService;
    private final NotificationHookConsumer notificationHookConsumer;
    private final Semaphore                asyncImportSemaphore;
    private final Configuration            applicationProperties;

    @Inject
    public ImportTaskListenerImpl(AsyncImportService asyncImportService, NotificationHookConsumer notificationHookConsumer) throws AtlasException {
        this(asyncImportService, notificationHookConsumer, new LinkedBlockingQueue<>());
    }

    public ImportTaskListenerImpl(AsyncImportService asyncImportService, NotificationHookConsumer notificationHookConsumer, BlockingQueue<String> requestQueue) throws AtlasException {
        this.asyncImportService       = asyncImportService;
        this.notificationHookConsumer = notificationHookConsumer;
        this.requestQueue             = requestQueue;
        this.asyncImportSemaphore     = new Semaphore(ASYNC_IMPORT_PERMITS);
        this.applicationProperties    = ApplicationProperties.get();
    }

    @Override
    public void start() throws AtlasException {
        if (HAConfiguration.isHAEnabled(applicationProperties)) {
            LOG.info("HA is enabled, not starting import consumers inline.");

            return;
        }

        startInternal();
    }

    @Override
    public void stop() throws AtlasException {
        try {
            stopImport();
        } finally {
            releaseAsyncImportSemaphore();
        }
    }

    @Override
    public void instanceIsActive() {
        LOG.info("Reacting to active state: initializing Kafka consumers");

        isActiveInstance = true;
        startInternal();
    }

    @Override
    public void instanceIsPassive() {
        isActiveInstance = false;
        try {
            stopImport();
        } finally {
            releaseAsyncImportSemaphore();
        }
    }

    @Override
    public int getHandlerOrder() {
        return ActiveStateChangeHandler.HandlerOrder.IMPORT_TASK_LISTENER.getOrder();
    }

    @Override
    public void onReceiveImportRequest(AtlasAsyncImportRequest importRequest) throws AtlasBaseException {
        try {
            LOG.info("==> onReceiveImportRequest(atlasAsyncImportRequest={})", importRequest);

            importRequest.setStatus(ImportStatus.WAITING);

            asyncImportService.updateImportRequest(importRequest);
            requestQueue.put(importRequest.getImportId());

            startNextImportInQueue();
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();

            LOG.warn("Failed to add import request: {} to the queue", importRequest.getImportId());

            throw new AtlasBaseException(IMPORT_QUEUEING_FAILED, e, importRequest.getImportId());
        } finally {
            LOG.info("<== onReceiveImportRequest(atlasAsyncImportRequest={})", importRequest);
        }
    }

    @Override
    public void onCompleteImportRequest(String importId) {
        LOG.info("==> onCompleteImportRequest(importId={})", importId);

        try {
            notificationHookConsumer.closeImportConsumer(importId, ASYNC_IMPORT_TOPIC_PREFIX.getString() + importId);
        } finally {
            releaseAsyncImportSemaphore();
            startNextImportInQueue();

            LOG.info("<== onCompleteImportRequest(importId={})", importId);
        }
    }

    @PreDestroy
    public void stopImport() {
        LOG.info("Shutting down import processor...");

        if (executorService == null) {
            LOG.info("Executor service is already null, nothing to shut down.");
            return;
        }
        executorService.shutdown(); // Initiate an orderly shutdown

        try {
            if (!executorService.awaitTermination(30, TimeUnit.SECONDS)) {
                LOG.warn("Executor service did not terminate gracefully within the timeout. Waiting longer...");

                // Retry shutdown before forcing it
                if (!executorService.awaitTermination(10, TimeUnit.SECONDS)) {
                    LOG.warn("Forcing shutdown...");

                    executorService.shutdownNow();
                }
            }
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();

            LOG.error("Shutdown interrupted. Forcing shutdown...");

            executorService.shutdownNow();
        }

        LOG.info("Import processor stopped.");
    }

    @VisibleForTesting
    void startInternal() {
        populateRequestQueue();

        if (!requestQueue.isEmpty()) {
            CompletableFuture.runAsync(this::startNextImportInQueue)
                    .exceptionally(ex -> {
                        LOG.error("Failed to start next import in queue", ex);

                        return null;
                    });
        }
    }

    @VisibleForTesting
    void startNextImportInQueue() {
        LOG.info("==> startNextImportInQueue()");

        startAsyncImportIfAvailable(null);

        LOG.info("<== startNextImportInQueue()");
    }

    @VisibleForTesting
    void startAsyncImportIfAvailable(String importId) {
        LOG.info("==> startAsyncImportIfAvailable()");

        if (!isActiveInstance) {
            LOG.warn("Import processing attempted while instance is passive. Skipping import.");
            return;
        }
        try {
            if (!asyncImportSemaphore.tryAcquire()) {
                LOG.info("An async import is in progress, import request is queued");

                return;
            }

            AtlasAsyncImportRequest nextImport = (importId != null) ? asyncImportService.fetchImportRequestByImportId(importId) : getNextImportFromQueue();

            if (isNotValidImportRequest(nextImport)) {
                releaseAsyncImportSemaphore();

                return;
            }

            ExecutorService exec = ensureExecutorAlive();
            if (exec != null) {
                exec.submit(() -> startImportConsumer(nextImport));
            } else {
                LOG.warn("No executor available to process import task (instance is passive).");
            }
        } catch (Exception e) {
            LOG.error("Error while starting the next import, releasing the lock if held", e);

            releaseAsyncImportSemaphore();
        } finally {
            LOG.info("<== startAsyncImportIfAvailable()");
        }
    }

    @VisibleForTesting
    AtlasAsyncImportRequest getNextImportFromQueue() {
        LOG.info("==> getNextImportFromQueue()");

        final int maxRetries = 5;
        int       retryCount = 0;

        while (retryCount < maxRetries) {
            try {
                String importId = requestQueue.poll(10, TimeUnit.SECONDS);

                if (importId == null) {
                    retryCount++;

                    LOG.warn("Still waiting for import request... (attempt {} of {})", retryCount, maxRetries);

                    continue;
                }

                // Reset retry count because we got a valid importId (even if it's invalid later)
                retryCount = 0;

                AtlasAsyncImportRequest importRequest = asyncImportService.fetchImportRequestByImportId(importId);

                if (isNotValidImportRequest(importRequest)) {
                    LOG.info("Import request {}, is not in a valid status to start import, hence skipping..", importRequest);

                    continue;
                }

                LOG.info("<== getImportIdFromQueue(nextImportId={})", importRequest.getImportId());

                return importRequest;
            } catch (InterruptedException e) {
                LOG.error("Thread interrupted while waiting for importId from the queue", e);

                // Restore the interrupt flag
                Thread.currentThread().interrupt();

                return null;
            }
        }

        LOG.error("Exceeded max retry attempts. Exiting...");

        return null;
    }

    @VisibleForTesting
    boolean isNotValidImportRequest(AtlasAsyncImportRequest importRequest) {
        return importRequest == null ||
                (!ImportStatus.WAITING.equals(importRequest.getStatus()) && !ImportStatus.PROCESSING.equals(importRequest.getStatus()));
    }

    @VisibleForTesting
    ExecutorService ensureExecutorAlive() {
        if (!isActiveInstance) {
            LOG.warn("Attempted to create executor while instance is passive. No executor will be created.");
            return null;
        }
        if (executorService == null || executorService.isShutdown() || executorService.isTerminated()) {
            synchronized (this) {
                if (executorService == null || executorService.isShutdown() || executorService.isTerminated()) {
                    executorService = Executors.newSingleThreadExecutor(new ThreadFactoryBuilder().setNameFormat(THREADNAME_PREFIX + " thread-%d")
                            .setUncaughtExceptionHandler((thread, throwable) -> LOG.error("Uncaught exception in thread {}: {}", thread.getName(), throwable.getMessage(), throwable)).build());
                    LOG.info("ExecutorService was recreated.");
                }
            }
        }
        return executorService;
    }

    void populateRequestQueue() {
        LOG.info("==> populateRequestQueue()");

        List<String> queuedImports     = asyncImportService.fetchQueuedImportRequests();
        List<String> inProgressImports = asyncImportService.fetchInProgressImportIds();

        if (queuedImports.isEmpty() && inProgressImports.isEmpty()) {
            LOG.info("populateRequestQueue(): no queued asynchronous import requests found.");
        } else {
            LOG.info("populateRequestQueue(): loaded {} asynchronous import requests (in-progress={}, queued={})", (inProgressImports.size() + queuedImports.size()), inProgressImports.size(), queuedImports.size());

            Stream.concat(inProgressImports.stream(), queuedImports.stream()).forEach(this::enqueueImportId);
        }

        LOG.info("<== populateRequestQueue()");
    }

    private void startImportConsumer(AtlasAsyncImportRequest importRequest) {
        try {
            LOG.info("==> startImportConsumer(atlasAsyncImportRequest={})", importRequest);

            notificationHookConsumer.startAsyncImportConsumer(NotificationInterface.NotificationType.ASYNC_IMPORT, importRequest.getImportId(), importRequest.getTopicName());

            importRequest.setStatus(ImportStatus.PROCESSING);
            importRequest.setProcessingStartTime(System.currentTimeMillis());
        } catch (Exception e) {
            LOG.error("Failed to start consumer for import: {}, marking import as failed", importRequest, e);

            importRequest.setStatus(ImportStatus.FAILED);
        } finally {
            asyncImportService.updateImportRequest(importRequest);

            if (ObjectUtils.equals(importRequest.getStatus(), ImportStatus.FAILED)) {
                onCompleteImportRequest(importRequest.getImportId());
            }

            LOG.info("<== startImportConsumer(atlasAsyncImportRequest={})", importRequest);
        }
    }

    private void releaseAsyncImportSemaphore() {
        LOG.info("==> releaseAsyncImportSemaphore()");

        if (asyncImportSemaphore.availablePermits() == 0) {
            asyncImportSemaphore.release();

            LOG.info("<== releaseAsyncImportSemaphore()");
        } else {
            LOG.info("<== releaseAsyncImportSemaphore(); no lock held");
        }
    }

    private void enqueueImportId(String importId) {
        try {
            if (!requestQueue.offer(importId, 5, TimeUnit.SECONDS)) {
                LOG.warn("populateRequestQueue(): failed to add import {} to the queue - enqueue timed out", importId);
            }
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();

            LOG.error("populateRequestQueue(): Failed to add import {} to the queue", importId, e);
        }
    }
}
