/*
 * 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.nifi.engine;

import org.apache.nifi.nar.NarThreadContextClassLoader;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.concurrent.Callable;
import java.util.concurrent.CancellationException;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.FutureTask;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

public final class FlowEngine extends ScheduledThreadPoolExecutor {

    private static final Logger logger = LoggerFactory.getLogger(FlowEngine.class);

    /**
     * Creates a new instance of FlowEngine
     *
     * @param corePoolSize the maximum number of threads available to tasks running in the engine.
     * @param threadNamePrefix for naming the thread
     */
    public FlowEngine(int corePoolSize, final String threadNamePrefix) {
        this(corePoolSize, threadNamePrefix, false);
    }

    /**
     * Creates a new instance of FlowEngine
     *
     * @param corePoolSize the maximum number of threads available to tasks running in the engine.
     * @param threadNamePrefix for thread naming
     * @param daemon if true, the thread pool will be populated with daemon threads, otherwise the threads will not be marked as daemon.
     */
    public FlowEngine(int corePoolSize, final String threadNamePrefix, final boolean daemon) {
        super(corePoolSize);

        final AtomicInteger threadIndex = new AtomicInteger(0);
        final ThreadFactory defaultThreadFactory = getThreadFactory();
        setThreadFactory(r -> {
            final Thread t = defaultThreadFactory.newThread(r);
            if (daemon) {
                t.setDaemon(true);
            }
            t.setName(threadNamePrefix + " Thread-" + threadIndex.incrementAndGet());
            return t;
        });
    }

    /**
     * Hook method called by the running thread whenever a runnable task is given to the thread to run.
     *
     * @param thread thread
     * @param runnable runnable
     */
    @Override
    protected void beforeExecute(final Thread thread, final Runnable runnable) {
        // Ensure classloader is correct
        thread.setContextClassLoader(NarThreadContextClassLoader.getInstance());
        super.beforeExecute(thread, runnable);
    }

    @Override
    public ScheduledFuture<?> schedule(final Runnable command, final long delay, final TimeUnit unit) {
        return super.schedule(wrap(command), delay, unit);
    }

    @Override
    public ScheduledFuture<?> scheduleAtFixedRate(final Runnable command, final long initialDelay, final long period, final TimeUnit unit) {
        return super.scheduleAtFixedRate(wrap(command), initialDelay, period, unit);
    }

    @Override
    public ScheduledFuture<?> scheduleWithFixedDelay(final Runnable command, final long initialDelay, final long delay, final TimeUnit unit) {
        return super.scheduleWithFixedDelay(wrap(command), initialDelay, delay, unit);
    }

    @Override
    public <V> ScheduledFuture<V> schedule(final Callable<V> callable, final long delay, final TimeUnit unit) {
        return super.schedule(wrap(callable), delay, unit);
    }

    private Runnable wrap(final Runnable runnable) {
        return () -> {
            try {
                runnable.run();
            } catch (final Throwable t) {
                logger.error("Uncaught Exception in Runnable task", t);
            }
        };
    }

    private <T> Callable<T> wrap(final Callable<T> callable) {
        return () -> {
            try {
                return callable.call();
            } catch (final Throwable t) {
                logger.error("Uncaught Exception in Callable task", t);
                throw t;
            }
        };
    }

    /**
     * Hook method called by the thread that executed the given runnable after execution of the runnable completed. Logs the fact of completion and any errors that might have occurred.
     *
     * @param runnable runnable
     * @param throwable throwable
     */
    @Override
    protected void afterExecute(final Runnable runnable, final Throwable throwable) {
        super.afterExecute(runnable, throwable);
        if (runnable instanceof FutureTask<?>) {
            final FutureTask<?> task = (FutureTask<?>) runnable;
            try {
                if (task.isDone()) {
                    if (task.isCancelled()) {
                        if (logger.isDebugEnabled()) {
                            logger.debug("A flow controller execution task '{}' has been cancelled.", runnable);
                        }
                    } else {
                        task.get(); //to raise any exceptions that might have occurred.
                        logger.debug("A Flow Controller execution task '{}' has completed.", runnable);
                    }
                }
            } catch (final CancellationException ce) {
                if (logger.isDebugEnabled()) {
                    logger.debug("A flow controller execution task '{}' has been cancelled.", runnable);
                }
            } catch (final InterruptedException ie) {
                if (logger.isDebugEnabled()) {
                    logger.debug("A flow controller execution task has been interrupted.", ie);
                }
            } catch (final ExecutionException ee) {
                logger.error("A flow controller task execution stopped abnormally", ee);
            }
        } else {
            logger.debug("A flow controller execution task '{}' has finished.", runnable);
        }
    }

    /**
     * Hook method called whenever the engine is terminated.
     */
    @Override
    protected void terminated() {
        super.terminated();
    }
}
