/*
 * 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.reef.runtime.common.evaluator.context;

import com.google.protobuf.ByteString;
import org.apache.reef.annotations.Provided;
import org.apache.reef.annotations.audience.EvaluatorSide;
import org.apache.reef.annotations.audience.Private;
import org.apache.reef.proto.EvaluatorRuntimeProtocol;
import org.apache.reef.proto.ReefServiceProtos;
import org.apache.reef.runtime.common.evaluator.HeartBeatManager;
import org.apache.reef.runtime.common.evaluator.task.TaskClientCodeException;
import org.apache.reef.runtime.common.utils.ExceptionCodec;
import org.apache.reef.tang.Configuration;
import org.apache.reef.tang.InjectionFuture;
import org.apache.reef.tang.exceptions.BindException;
import org.apache.reef.tang.formats.ConfigurationSerializer;
import org.apache.reef.util.Optional;

import javax.inject.Inject;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Stack;
import java.util.logging.Level;
import java.util.logging.Logger;

/**
 * Manages the stack of context in the Evaluator.
 */
@Private
@EvaluatorSide
@Provided
public final class ContextManager implements AutoCloseable {

  private static final Logger LOG = Logger.getLogger(ContextManager.class.getName());

  /**
   * The stack of context.
   */
  private final Stack<ContextRuntime> contextStack = new Stack<>();

  /**
   * Used to instantiate the root context.
   */
  private final InjectionFuture<RootContextLauncher> launchContext;

  /**
   * Used for status reporting to the Driver.
   */
  private final HeartBeatManager heartBeatManager;

  /**
   * To serialize Configurations.
   */
  private final ConfigurationSerializer configurationSerializer;

  private final ExceptionCodec exceptionCodec;

  /**
   * @param launchContext           to instantiate the root context.
   * @param heartBeatManager        for status reporting to the Driver.
   * @param configurationSerializer
   * @param exceptionCodec
   */
  @Inject
  ContextManager(final InjectionFuture<RootContextLauncher> launchContext,
                 final HeartBeatManager heartBeatManager,
                 final ConfigurationSerializer configurationSerializer,
                 final ExceptionCodec exceptionCodec) {
    this.launchContext = launchContext;
    this.heartBeatManager = heartBeatManager;
    this.configurationSerializer = configurationSerializer;
    this.exceptionCodec = exceptionCodec;
  }

  /**
   * Start the context manager. This initiates the root context.
   *
   * @throws ContextClientCodeException if the root context can't be instantiated.
   */
  public void start() throws ContextClientCodeException {
    synchronized (this.contextStack) {
      LOG.log(Level.FINEST, "Instantiating root context.");
      this.contextStack.push(this.launchContext.get().getRootContext());

      if (this.launchContext.get().getInitialTaskConfiguration().isPresent()) {
        LOG.log(Level.FINEST, "Launching the initial Task");
        try {
          this.contextStack.peek().startTask(this.launchContext.get().getInitialTaskConfiguration().get());
        } catch (final TaskClientCodeException e) {
          this.handleTaskException(e);
        }
      }
    }
  }

  /**
   * Shuts down. This forecefully kills the Task if there is one and then shuts down all Contexts on the stack,
   * starting at the top.
   */
  @Override
  public void close() {
    synchronized (this.contextStack) {
      if (!this.contextStackIsEmpty()) {
        this.contextStack.lastElement().close();
      }
    }
  }

  /**
   * @return true if there is no context.
   */
  public boolean contextStackIsEmpty() {
    synchronized (this.contextStack) {
      return this.contextStack.isEmpty();
    }
  }

  /**
   * Processes the given ContextControlProto to launch / close / suspend Tasks and Contexts.
   * <p>
   * This also triggers the HeartBeatManager to send a heartbeat with the result of this operation.
   *
   * @param controlMessage the message to process
   */
  public void handleContextControlProtocol(
      final EvaluatorRuntimeProtocol.ContextControlProto controlMessage) {

    synchronized (this.heartBeatManager) {
      try {
        if (controlMessage.hasAddContext() && controlMessage.hasRemoveContext()) {
          throw new IllegalArgumentException(
              "Received a message with both add and remove context. This is unsupported.");
        }

        final byte[] message = controlMessage.hasTaskMessage() ?
            controlMessage.getTaskMessage().toByteArray() : null;

        if (controlMessage.hasAddContext()) {
          this.addContext(controlMessage.getAddContext());
          if (controlMessage.hasStartTask()) {
            // We support submitContextAndTask()
            this.startTask(controlMessage.getStartTask());
          } else {
            // We need to trigger a heartbeat here.
            // In other cases, the heartbeat will be triggered by the TaskRuntime
            // Therefore this call can not go into addContext.
            this.heartBeatManager.sendHeartbeat();
          }
        } else if (controlMessage.hasRemoveContext()) {
          this.removeContext(controlMessage.getRemoveContext().getContextId());
        } else if (controlMessage.hasStartTask()) {
          this.startTask(controlMessage.getStartTask());
        } else if (controlMessage.hasStopTask()) {
          this.contextStack.peek().closeTask(message);
        } else if (controlMessage.hasSuspendTask()) {
          this.contextStack.peek().suspendTask(message);
        } else if (controlMessage.hasTaskMessage()) {
          this.contextStack.peek().deliverTaskMessage(message);
        } else if (controlMessage.hasContextMessage()) {
          final EvaluatorRuntimeProtocol.ContextMessageProto contextMessageProto = controlMessage.getContextMessage();
          boolean deliveredMessage = false;
          for (final ContextRuntime context : this.contextStack) {
            if (context.getIdentifier().equals(contextMessageProto.getContextId())) {
              context.handleContextMessage(contextMessageProto.getMessage().toByteArray());
              deliveredMessage = true;
              break;
            }
          }
          if (!deliveredMessage) {
            throw new IllegalStateException(
                "Sent message to unknown context " + contextMessageProto.getContextId());
          }
        } else {
          throw new RuntimeException("Unknown task control message: " + controlMessage);
        }
      } catch (final TaskClientCodeException e) {
        this.handleTaskException(e);
      } catch (final ContextClientCodeException e) {
        this.handleContextException(e);
      }
    }

  }

  /**
   * @return the TaskStatusProto of the currently running task, if there is any
   */
  public Optional<ReefServiceProtos.TaskStatusProto> getTaskStatus() {
    synchronized (this.contextStack) {
      if (this.contextStack.isEmpty()) {
        throw new RuntimeException(
            "Asked for a Task status while there isn't even a context running.");
      }
      return this.contextStack.peek().getTaskStatus();
    }
  }

  /**
   * @return the status of all context in the stack.
   */
  public Collection<ReefServiceProtos.ContextStatusProto> getContextStatusCollection() {
    synchronized (this.contextStack) {
      final List<ReefServiceProtos.ContextStatusProto> result = new ArrayList<>(this.contextStack.size());
      for (final ContextRuntime contextRuntime : this.contextStack) {
        final ReefServiceProtos.ContextStatusProto contextStatusProto = contextRuntime.getContextStatus();
        LOG.log(Level.FINEST, "Add context status: {0}", contextStatusProto);
        result.add(contextStatusProto);
      }
      return result;
    }
  }

  /**
   * Add a context to the stack.
   *
   * @param addContextProto
   * @throws ContextClientCodeException if there is a client code related issue.
   */
  private void addContext(
      final EvaluatorRuntimeProtocol.AddContextProto addContextProto)
      throws ContextClientCodeException {

    synchronized (this.contextStack) {
      try {

        final ContextRuntime currentTopContext = this.contextStack.peek();

        if (!currentTopContext.getIdentifier().equals(addContextProto.getParentContextId())) {
          throw new IllegalStateException("Trying to instantiate a child context on context with id `" +
              addContextProto.getParentContextId() + "` while the current top context id is `" +
              currentTopContext.getIdentifier() + "`");
        }

        final Configuration contextConfiguration =
            this.configurationSerializer.fromString(addContextProto.getContextConfiguration());

        final ContextRuntime newTopContext;
        if (addContextProto.hasServiceConfiguration()) {
          newTopContext = currentTopContext.spawnChildContext(contextConfiguration,
              this.configurationSerializer.fromString(addContextProto.getServiceConfiguration()));
        } else {
          newTopContext = currentTopContext.spawnChildContext(contextConfiguration);
        }

        this.contextStack.push(newTopContext);

      } catch (final IOException | BindException e) {
        throw new RuntimeException("Unable to read configuration.", e);
      }

    }
  }

  /**
   * Remove the context with the given ID from the stack.
   *
   * @throws IllegalStateException if the given ID does not refer to the top of stack.
   */
  private void removeContext(final String contextID) {

    synchronized (this.contextStack) {

      if (!contextID.equals(this.contextStack.peek().getIdentifier())) {
        throw new IllegalStateException("Trying to close context with id `" + contextID +
            "`. But the top context has id `" +
            this.contextStack.peek().getIdentifier() + "`");
      }

      this.contextStack.peek().close();
      if (this.contextStack.size() > 1) {
        /* We did not close the root context. Therefore, we need to inform the
         * driver explicitly that this context is closed. The root context notification
         * is implicit in the Evaluator close/done notification.
         */
        this.heartBeatManager.sendHeartbeat(); // Ensure Driver gets notified of context DONE state
      }
      this.contextStack.pop();
      /*
      * At this moment, the Evaluator is actually idle and has some time till the Driver sends it additional work.
      * Also, a potentially large object graph just became orphaned: all the objects instantiated by the context
      * and service injectors can now be garbage collected. So GC call is justified.
      * */
      System.gc();
    }
  }

  /**
   * Launch a Task.
   */
  private void startTask(
      final EvaluatorRuntimeProtocol.StartTaskProto startTaskProto) throws TaskClientCodeException {

    synchronized (this.contextStack) {

      final ContextRuntime currentActiveContext = this.contextStack.peek();

      final String expectedContextId = startTaskProto.getContextId();
      if (!expectedContextId.equals(currentActiveContext.getIdentifier())) {
        throw new IllegalStateException("Task expected context `" + expectedContextId +
            "` but the active context has ID `" + currentActiveContext.getIdentifier() + "`");
      }

      try {
        final Configuration taskConfig =
            this.configurationSerializer.fromString(startTaskProto.getConfiguration());
        currentActiveContext.startTask(taskConfig);
      } catch (IOException | BindException e) {
        throw new RuntimeException("Unable to read configuration.", e);
      }
    }
  }

  /**
   * THIS ASSUMES THAT IT IS CALLED ON A THREAD HOLDING THE LOCK ON THE HeartBeatManager.
   */
  private void handleTaskException(final TaskClientCodeException e) {

    LOG.log(Level.SEVERE, "TaskClientCodeException", e);

    final ByteString exception = ByteString.copyFrom(this.exceptionCodec.toBytes(e.getCause()));

    final ReefServiceProtos.TaskStatusProto taskStatus =
        ReefServiceProtos.TaskStatusProto.newBuilder()
            .setContextId(e.getContextId())
            .setTaskId(e.getTaskId())
            .setResult(exception)
            .setState(ReefServiceProtos.State.FAILED)
            .build();

    LOG.log(Level.SEVERE, "Sending heartbeat: {0}", taskStatus);

    this.heartBeatManager.sendTaskStatus(taskStatus);
  }

  /**
   * THIS ASSUMES THAT IT IS CALLED ON A THREAD HOLDING THE LOCK ON THE HeartBeatManager.
   */
  private void handleContextException(final ContextClientCodeException e) {

    LOG.log(Level.SEVERE, "ContextClientCodeException", e);

    final ByteString exception = ByteString.copyFrom(this.exceptionCodec.toBytes(e.getCause()));

    final ReefServiceProtos.ContextStatusProto.Builder contextStatusBuilder =
        ReefServiceProtos.ContextStatusProto.newBuilder()
            .setContextId(e.getContextID())
            .setContextState(ReefServiceProtos.ContextStatusProto.State.FAIL)
            .setError(exception);

    if (e.getParentID().isPresent()) {
      contextStatusBuilder.setParentId(e.getParentID().get());
    }

    final ReefServiceProtos.ContextStatusProto contextStatus = contextStatusBuilder.build();

    LOG.log(Level.SEVERE, "Sending heartbeat: {0}", contextStatus);

    this.heartBeatManager.sendContextStatus(contextStatus);
  }
}