/*
 * 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.kafka.trogdor.agent;

import org.apache.kafka.common.KafkaException;
import org.apache.kafka.common.KafkaFuture;
import org.apache.kafka.common.internals.KafkaFutureImpl;
import org.apache.kafka.common.utils.Scheduler;
import org.apache.kafka.common.utils.ThreadUtils;
import org.apache.kafka.common.utils.Time;
import org.apache.kafka.common.utils.Utils;
import org.apache.kafka.trogdor.common.Platform;
import org.apache.kafka.trogdor.rest.RequestConflictException;
import org.apache.kafka.trogdor.rest.WorkerDone;
import org.apache.kafka.trogdor.rest.WorkerRunning;
import org.apache.kafka.trogdor.rest.WorkerStarting;
import org.apache.kafka.trogdor.rest.WorkerState;
import org.apache.kafka.trogdor.rest.WorkerStopping;
import org.apache.kafka.trogdor.task.AgentWorkerStatusTracker;
import org.apache.kafka.trogdor.task.TaskSpec;
import org.apache.kafka.trogdor.task.TaskWorker;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.Map;
import java.util.TreeMap;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

public final class WorkerManager {
    private static final Logger log = LoggerFactory.getLogger(WorkerManager.class);

    /**
     * The platform to use.
     */
    private final Platform platform;

    /**
     * The name of this node.
     */
    private final String nodeName;

    /**
     * The scheduler to use.
     */
    private final Scheduler scheduler;

    /**
     * The clock to use.
     */
    private final Time time;

    /**
     * A map of task IDs to Work objects.
     */
    private final Map<Long, Worker> workers;

    /**
     * An ExecutorService used to schedule events in the future.
     */
    private final ScheduledExecutorService stateChangeExecutor;

    /**
     * An ExecutorService used to clean up TaskWorkers.
     */
    private final ExecutorService workerCleanupExecutor;

    /**
     * An ExecutorService to help with shutting down.
     */
    private final ScheduledExecutorService shutdownExecutor;

    /**
     * The shutdown manager.
     */
    private final ShutdownManager shutdownManager = new ShutdownManager();

    /**
     * The shutdown manager handles shutting down gracefully.
     *
     * We can shut down gracefully only when all the references handed out
     * by the ShutdownManager has been closed, and the shutdown bit has
     * been set.  RPC operations hold a reference for the duration of their
     * execution, and so do Workers which have not been shut down.
     * This prevents us from shutting down in the middle of an RPC, or with
     * workers which are still running.
     */
    static class ShutdownManager {
        private boolean shutdown = false;
        private long refCount = 0;

        class Reference implements AutoCloseable {
            AtomicBoolean closed = new AtomicBoolean(false);

            @Override
            public void close() {
                if (closed.compareAndSet(false, true)) {
                    synchronized (ShutdownManager.this) {
                        refCount--;
                        if (shutdown && (refCount == 0)) {
                            ShutdownManager.this.notifyAll();
                        }
                    }
                }
            }
        }

        synchronized Reference takeReference() {
            if (shutdown) {
                throw new KafkaException("WorkerManager is shut down.");
            }
            refCount++;
            return new Reference();
        }

        synchronized boolean shutdown() {
            if (shutdown) {
                return false;
            }
            shutdown = true;
            if (refCount == 0) {
                this.notifyAll();
            }
            return true;
        }

        synchronized void waitForQuiescence() throws InterruptedException {
            while ((!shutdown) || (refCount > 0)) {
                this.wait();
            }
        }
    }

    WorkerManager(Platform platform, Scheduler scheduler) {
        this.platform = platform;
        this.nodeName = platform.curNode().name();
        this.scheduler = scheduler;
        this.time = scheduler.time();
        this.workers = new HashMap<>();
        this.stateChangeExecutor = Executors.newSingleThreadScheduledExecutor(
                ThreadUtils.createThreadFactory("WorkerManagerStateThread", false));
        this.workerCleanupExecutor = Executors.newCachedThreadPool(
            ThreadUtils.createThreadFactory("WorkerCleanupThread%d", false));
        this.shutdownExecutor = Executors.newScheduledThreadPool(0,
            ThreadUtils.createThreadFactory("WorkerManagerShutdownThread%d", false));
    }

    enum State {
        STARTING,
        CANCELLING,
        RUNNING,
        STOPPING,
        DONE,
    }

    /**
     * A worker which is being tracked.
     */
    class Worker {
        /**
         * The worker ID.
         */
        private final long workerId;

        /**
         * The task ID.
         */
        private final String taskId;

        /**
         * The task specification.
         */
        private final TaskSpec spec;

        /**
         * The work which this worker is performing.
         */
        private final TaskWorker taskWorker;

        /**
         * The worker status.
         */
        private final AgentWorkerStatusTracker status = new AgentWorkerStatusTracker();

        /**
         * The time when this task was started.
         */
        private final long startedMs;

        /**
         * The work state.
         */
        private State state = State.STARTING;

        /**
         * The time when this task was completed, or -1 if it has not been.
         */
        private long doneMs = -1;

        /**
         * The worker error.
         */
        private String error = "";

        /**
         * If there is a task timeout scheduled, this is a future which can
         * be used to cancel it.
         */
        private Future<Void> timeoutFuture = null;

        /**
         * A future which is completed when the task transitions to DONE state.
         */
        private KafkaFutureImpl<String> doneFuture = null;

        /**
         * A shutdown manager reference which will keep the WorkerManager
         * alive for as long as this worker is alive.
         */
        private ShutdownManager.Reference reference;

        /**
         * Whether we should destroy the records of this worker once it stops.
         */
        private boolean mustDestroy = false;

        Worker(long workerId, String taskId, TaskSpec spec, long now) {
            this.workerId = workerId;
            this.taskId = taskId;
            this.spec = spec;
            this.taskWorker = spec.newTaskWorker(taskId);
            this.startedMs = now;
            this.reference = shutdownManager.takeReference();
        }

        long workerId() {
            return workerId;
        }

        String taskId() {
            return taskId;
        }

        TaskSpec spec() {
            return spec;
        }

        WorkerState state() {
            switch (state) {
                case STARTING:
                    return new WorkerStarting(taskId, spec);
                case RUNNING:
                    return new WorkerRunning(taskId, spec, startedMs, status.get());
                case CANCELLING:
                case STOPPING:
                    return new WorkerStopping(taskId, spec, startedMs, status.get());
                case DONE:
                    return new WorkerDone(taskId, spec, startedMs, doneMs, status.get(), error);
            }
            throw new RuntimeException("unreachable");
        }

        void transitionToRunning() {
            state = State.RUNNING;
            timeoutFuture = scheduler.schedule(stateChangeExecutor,
                new StopWorker(workerId, false),
                Math.max(0, spec.endMs() - time.milliseconds()));
        }

        Future<Void> transitionToStopping() {
            state = State.STOPPING;
            if (timeoutFuture != null) {
                timeoutFuture.cancel(false);
                timeoutFuture = null;
            }
            return workerCleanupExecutor.submit(new HaltWorker(this));
        }

        void transitionToDone() {
            state = State.DONE;
            doneMs = time.milliseconds();
            if (reference != null) {
                reference.close();
                reference = null;
            }
            doneFuture.complete(error);
        }

        @Override
        public String toString() {
            return String.format("%s_%d", taskId, workerId);
        }
    }

    public KafkaFuture<String> createWorker(long workerId, String taskId, TaskSpec spec) throws Throwable {
        try (ShutdownManager.Reference ref = shutdownManager.takeReference()) {
            final Worker worker = stateChangeExecutor.
                submit(new CreateWorker(workerId, taskId, spec, time.milliseconds())).get();
            if (worker.doneFuture != null) {
                log.info("{}: Ignoring request to create worker {}, because there is already " +
                    "a worker with that id.", nodeName, workerId);
                return worker.doneFuture;
            }
            worker.doneFuture = new KafkaFutureImpl<>();
            if (worker.spec.endMs() <= time.milliseconds()) {
                log.info("{}: Will not run worker {} as it has expired.", nodeName, worker);
                stateChangeExecutor.submit(new HandleWorkerHalting(worker,
                    "worker expired", true));
                return worker.doneFuture;
            }
            KafkaFutureImpl<String> haltFuture = new KafkaFutureImpl<>();
            haltFuture.thenApply(errorString -> {
                if (errorString == null)
                    errorString = "";
                if (errorString.isEmpty()) {
                    log.info("{}: Worker {} is halting.", nodeName, worker);
                } else {
                    log.info("{}: Worker {} is halting with error {}",
                        nodeName, worker, errorString);
                }
                stateChangeExecutor.submit(
                    new HandleWorkerHalting(worker, errorString, false));
                return null;
            });
            try {
                worker.taskWorker.start(platform, worker.status, haltFuture);
            } catch (Exception e) {
                log.info("{}: Worker {} start() exception", nodeName, worker, e);
                stateChangeExecutor.submit(new HandleWorkerHalting(worker,
                    "worker.start() exception: " + Utils.stackTrace(e), true));
            }
            stateChangeExecutor.submit(new FinishCreatingWorker(worker));
            return worker.doneFuture;
        } catch (ExecutionException e) {
            if (e.getCause() instanceof RequestConflictException) {
                log.info("{}: request conflict while creating worker {} for task {} with spec {}.",
                    nodeName, workerId, taskId, spec);
            } else {
                log.info("{}: Error creating worker {} for task {} with spec {}",
                    nodeName, workerId, taskId, spec, e);
            }
            throw e.getCause();
        }
    }

    /**
     * Handles a request to create a new worker.  Processed by the state change thread.
     */
    class CreateWorker implements Callable<Worker> {
        private final long workerId;
        private final String taskId;
        private final TaskSpec spec;
        private final long now;

        CreateWorker(long workerId, String taskId, TaskSpec spec, long now) {
            this.workerId = workerId;
            this.taskId = taskId;
            this.spec = spec;
            this.now = now;
        }

        @Override
        public Worker call() throws Exception {
            try {
                Worker worker = workers.get(workerId);
                if (worker != null) {
                    if (!worker.taskId().equals(taskId)) {
                        throw new RequestConflictException("There is already a worker ID " + workerId +
                            " with a different task ID.");
                    } else if (!worker.spec().equals(spec)) {
                        throw new RequestConflictException("There is already a worker ID " + workerId +
                            " with a different task spec.");
                    } else {
                        return worker;
                    }
                }
                worker = new Worker(workerId, taskId, spec, now);
                workers.put(workerId, worker);
                log.info("{}: Created worker {} with spec {}", nodeName, worker, spec);
                return worker;
            } catch (Exception e) {
                log.info("{}: unable to create worker {} for task {}, with spec {}",
                    nodeName, workerId, taskId, spec, e);
                throw e;
            }
        }
    }

    /**
     * Finish creating a Worker.  Processed by the state change thread.
     */
    class FinishCreatingWorker implements Callable<Void> {
        private final Worker worker;

        FinishCreatingWorker(Worker worker) {
            this.worker = worker;
        }

        @Override
        public Void call() throws Exception {
            switch (worker.state) {
                case CANCELLING:
                    log.info("{}: Worker {} was cancelled while it was starting up.  " +
                        "Transitioning to STOPPING.", nodeName, worker);
                    worker.transitionToStopping();
                    break;
                case STARTING:
                    log.info("{}: Worker {} is now RUNNING.  Scheduled to stop in {} ms.",
                        nodeName, worker, worker.spec.durationMs());
                    worker.transitionToRunning();
                    break;
                default:
                    break;
            }
            return null;
        }
    }

    /**
     * Handles a worker halting.  Processed by the state change thread.
     */
    class HandleWorkerHalting implements Callable<Void> {
        private final Worker worker;
        private final String failure;
        private final boolean startupHalt;

        HandleWorkerHalting(Worker worker, String failure, boolean startupHalt) {
            this.worker = worker;
            this.failure = failure;
            this.startupHalt = startupHalt;
        }

        @Override
        public Void call() throws Exception {
            if (worker.error.isEmpty()) {
                worker.error = failure;
            }
            String verb = (worker.error.isEmpty()) ? "halting" :
                "halting with error [" + worker.error + "]";
            switch (worker.state) {
                case STARTING:
                    if (startupHalt) {
                        log.info("{}: Worker {} {} during startup.  Transitioning to DONE.",
                            nodeName, worker, verb);
                        worker.transitionToDone();
                    } else {
                        log.info("{}: Worker {} {} during startup.  Transitioning to CANCELLING.",
                            nodeName, worker, verb);
                        worker.state = State.CANCELLING;
                    }
                    break;
                case CANCELLING:
                    log.info("{}: Cancelling worker {} {}.  ",
                            nodeName, worker, verb);
                    break;
                case RUNNING:
                    log.info("{}: Running worker {} {}.  Transitioning to STOPPING.",
                        nodeName, worker, verb);
                    worker.transitionToStopping();
                    break;
                case STOPPING:
                    log.info("{}: Stopping worker {} {}.", nodeName, worker, verb);
                    break;
                case DONE:
                    log.info("{}: Can't halt worker {} because it is already DONE.",
                        nodeName, worker);
                    break;
            }
            return null;
        }
    }

    /**
     * Transitions a worker to WorkerDone.  Processed by the state change thread.
     */
    class CompleteWorker implements Callable<Void> {
        private final Worker worker;

        private final String failure;

        CompleteWorker(Worker worker, String failure) {
            this.worker = worker;
            this.failure = failure;
        }

        @Override
        public Void call() throws Exception {
            if (worker.error.isEmpty() && !failure.isEmpty()) {
                worker.error = failure;
            }
            worker.transitionToDone();
            if (worker.mustDestroy) {
                log.info("{}: destroying worker {} with error {}",
                    nodeName, worker, worker.error);
                workers.remove(worker.workerId);
            } else {
                log.info("{}: completed worker {} with error {}",
                    nodeName, worker, worker.error);
            }
            return null;
        }
    }

    public void stopWorker(long workerId, boolean mustDestroy) throws Throwable {
        try (ShutdownManager.Reference ref = shutdownManager.takeReference()) {
            stateChangeExecutor.submit(new StopWorker(workerId, mustDestroy)).get();
        } catch (ExecutionException e) {
            throw e.getCause();
        }
    }

    /**
     * Stops a worker.  Processed by the state change thread.
     */
    class StopWorker implements Callable<Void> {
        private final long workerId;
        private final boolean mustDestroy;

        StopWorker(long workerId, boolean mustDestroy) {
            this.workerId = workerId;
            this.mustDestroy = mustDestroy;
        }

        @Override
        public Void call() throws Exception {
            Worker worker = workers.get(workerId);
            if (worker == null) {
                log.info("{}: Can't stop worker {} because there is no worker with that ID.",
                    nodeName, workerId);
                return null;
            }
            if (mustDestroy) {
                worker.mustDestroy = true;
            }
            switch (worker.state) {
                case STARTING:
                    log.info("{}: Cancelling worker {} during its startup process.",
                        nodeName, worker);
                    worker.state = State.CANCELLING;
                    break;
                case CANCELLING:
                    log.info("{}: Can't stop worker {}, because it is already being " +
                        "cancelled.", nodeName, worker);
                    break;
                case RUNNING:
                    log.info("{}: Stopping running worker {}.", nodeName, worker);
                    worker.transitionToStopping();
                    break;
                case STOPPING:
                    log.info("{}: Can't stop worker {}, because it is already " +
                            "stopping.", nodeName, worker);
                    break;
                case DONE:
                    if (worker.mustDestroy) {
                        log.info("{}: destroying worker {} with error {}",
                            nodeName, worker, worker.error);
                        workers.remove(worker.workerId);
                    } else {
                        log.debug("{}: Can't stop worker {}, because it is already done.",
                            nodeName, worker);
                    }
                    break;
            }
            return null;
        }
    }

    /**
     * Cleans up the resources associated with a worker.  Processed by the worker
     * cleanup thread pool.
     */
    class HaltWorker implements Callable<Void> {
        private final Worker worker;

        HaltWorker(Worker worker) {
            this.worker = worker;
        }

        @Override
        public Void call() throws Exception {
            String failure = "";
            try {
                worker.taskWorker.stop(platform);
            } catch (Exception exception) {
                log.error("{}: worker.stop() exception", nodeName, exception);
                failure = exception.getMessage();
            }
            stateChangeExecutor.submit(new CompleteWorker(worker, failure));
            return null;
        }
    }

    public TreeMap<Long, WorkerState> workerStates() throws Exception {
        try (ShutdownManager.Reference ref = shutdownManager.takeReference()) {
            return stateChangeExecutor.submit(new GetWorkerStates()).get();
        }
    }

    class GetWorkerStates implements Callable<TreeMap<Long, WorkerState>> {
        @Override
        public TreeMap<Long, WorkerState> call() throws Exception {
            TreeMap<Long, WorkerState> workerMap = new TreeMap<>();
            for (Worker worker : workers.values()) {
                workerMap.put(worker.workerId(), worker.state());
            }
            return workerMap;
        }
    }

    public void beginShutdown() throws Exception {
        if (shutdownManager.shutdown()) {
            shutdownExecutor.submit(new Shutdown());
        }
    }

    public void waitForShutdown() throws Exception {
        while (!shutdownExecutor.isShutdown()) {
            shutdownExecutor.awaitTermination(1, TimeUnit.DAYS);
        }
    }

    class Shutdown implements Callable<Void> {
        @Override
        public Void call() throws Exception {
            log.info("{}: Shutting down WorkerManager.", nodeName);
            try {
                stateChangeExecutor.submit(new DestroyAllWorkers()).get();
                log.info("{}: Waiting for shutdownManager quiescence...", nodeName);
                shutdownManager.waitForQuiescence();
                workerCleanupExecutor.shutdownNow();
                stateChangeExecutor.shutdownNow();
                log.info("{}: Waiting for workerCleanupExecutor to terminate...", nodeName);
                workerCleanupExecutor.awaitTermination(1, TimeUnit.DAYS);
                log.info("{}: Waiting for stateChangeExecutor to terminate...", nodeName);
                stateChangeExecutor.awaitTermination(1, TimeUnit.DAYS);
                log.info("{}: Shutting down shutdownExecutor.", nodeName);
                shutdownExecutor.shutdown();
            } catch (Exception e) {
                log.info("{}: Caught exception while shutting down WorkerManager", nodeName, e);
                throw e;
            }
            return null;
        }
    }

    /**
     * Begins the process of destroying all workers.  Processed by the state change thread.
     */
    class DestroyAllWorkers implements Callable<Void> {
        @Override
        public Void call() throws Exception {
            log.info("{}: Destroying all workers.", nodeName);

            // StopWorker may remove elements from the set of worker IDs.  That might generate
            // a ConcurrentModificationException if we were iterating over the worker ID
            // set directly.  Therefore, we make a copy of the worker IDs here and iterate
            // over that instead.
            //
            // Note that there is no possible way that more worker IDs can be added while this
            // callable is running, because the state change executor is single-threaded.
            ArrayList<Long> workerIds = new ArrayList<>(workers.keySet());

            for (long workerId : workerIds) {
                try {
                    new StopWorker(workerId, true).call();
                } catch (Exception e) {
                    log.error("Failed to stop worker {}", workerId, e);
                }
            }
            return null;
        }
    }

}
