package org.apache.helix.messaging.handling;

/*
 * 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.
 */

import java.util.Date;
import java.util.concurrent.ConcurrentHashMap;

import org.apache.helix.HelixDataAccessor;
import org.apache.helix.HelixManager;
import org.apache.helix.HelixRollbackException;
import org.apache.helix.InstanceType;
import org.apache.helix.NotificationContext;
import org.apache.helix.NotificationContext.MapKey;
import org.apache.helix.PropertyKey.Builder;
import org.apache.helix.messaging.handling.MessageHandler.ErrorCode;
import org.apache.helix.messaging.handling.MessageHandler.ErrorType;
import org.apache.helix.model.Message;
import org.apache.helix.model.Message.Attributes;
import org.apache.helix.model.Message.MessageType;
import org.apache.helix.monitoring.StateTransitionContext;
import org.apache.helix.monitoring.StateTransitionDataPoint;
import org.apache.helix.monitoring.mbeans.ParticipantMessageMonitor;
import org.apache.helix.util.StatusUpdateUtil;
import org.apache.log4j.Logger;

public class HelixTask implements MessageTask {
  private static Logger logger = Logger.getLogger(HelixTask.class);
  private final Message _message;
  private final MessageHandler _handler;
  private final NotificationContext _notificationContext;
  private final HelixManager _manager;
  StatusUpdateUtil _statusUpdateUtil;
  HelixTaskExecutor _executor;
  volatile boolean _isTimeout = false;
  volatile boolean _isStarted = false;
  volatile boolean _isCancelled = false;

  public HelixTask(Message message, NotificationContext notificationContext,
      MessageHandler handler, HelixTaskExecutor executor) {
    this._notificationContext = notificationContext;
    this._message = message;
    this._handler = handler;
    this._manager = notificationContext.getManager();
    _statusUpdateUtil = new StatusUpdateUtil();
    _executor = executor;
  }

  @Override
  public HelixTaskResult call() {
    HelixTaskResult taskResult = null;

    ErrorType type = null;
    ErrorCode code = null;

    long start = System.currentTimeMillis();
    logger.info("handling task: " + getTaskId() + " begin, at: " + start);
    HelixDataAccessor accessor = _manager.getHelixDataAccessor();
    _statusUpdateUtil.logInfo(_message, HelixTask.class, "Message handling task begin execute",
        accessor);
    _message.setExecuteStartTimeStamp(new Date().getTime());

    // add a concurrent map to hold currentStateUpdates for sub-messages of a batch-message
    // partitionName -> csUpdate
    if (_message.getBatchMessageMode() == true) {
      _notificationContext.add(MapKey.CURRENT_STATE_UPDATE.toString(),
          new ConcurrentHashMap<String, CurrentStateUpdate>());
    }

    // Handle the message
    try {
      setStarted();
      taskResult = _handler.handleMessage();
    } catch (InterruptedException e) {
      taskResult = new HelixTaskResult();
      taskResult.setException(e);
      taskResult.setInterrupted(true);

      _statusUpdateUtil.logError(_message, HelixTask.class, e,
          "State transition interrupted, timeout:" + _isTimeout, accessor);
      logger.info("Message " + _message.getMsgId() + " is interrupted");
    } catch (Exception e) {
      taskResult = new HelixTaskResult();
      taskResult.setException(e);
      taskResult.setMessage(e.getMessage());

      String errorMessage =
          "Exception while executing a message. " + e + " msgId: " + _message.getMsgId()
              + " type: " + _message.getMsgType();
      logger.error(errorMessage, e);
      _statusUpdateUtil.logError(_message, HelixTask.class, e, errorMessage, accessor);
      _executor.getParticipantMonitor().reportProcessedMessage(_message, ParticipantMessageMonitor.ProcessedMessageState.FAILED);
    }

    // cancel timeout task
    _executor.cancelTimeoutTask(this);

    Exception exception = null;
    try {
      if (taskResult.isSuccess()) {
        _statusUpdateUtil
            .logInfo(_message, _handler.getClass(), "Message handling task completed successfully", accessor);
        logger.info("Message " + _message.getMsgId() + " completed.");
        _executor.getParticipantMonitor().reportProcessedMessage(_message, ParticipantMessageMonitor.ProcessedMessageState.COMPLETED);
      } else {
        type = ErrorType.INTERNAL;

        if (taskResult.isInterrupted()) {
          logger.info("Message " + _message.getMsgId() + " is interrupted");
          code = _isTimeout ? ErrorCode.TIMEOUT : ErrorCode.CANCEL;
          if (_isTimeout) {
            int retryCount = _message.getRetryCount();
            logger.info("Message timeout, retry count: " + retryCount + " msgId:"
                + _message.getMsgId());
            _statusUpdateUtil.logInfo(_message, _handler.getClass(),
                "Message handling task timeout, retryCount:" + retryCount, accessor);
            // Notify the handler that timeout happens, and the number of retries left
            // In case timeout happens (time out and also interrupted)
            // we should retry the execution of the message by re-schedule it in
            if (retryCount > 0) {
              _message.setRetryCount(retryCount - 1);
              HelixTask task = new HelixTask(_message, _notificationContext, _handler, _executor);
              _executor.scheduleTask(task);
              return taskResult;
            }
          }
        } else if (taskResult.isCancelled()) {
          // Cancellation success, report message complete
          type = null;
          _statusUpdateUtil
              .logInfo(_message, _handler.getClass(), "Cancellation completed successfully",
                  accessor);
          _executor.getParticipantMonitor().reportProcessedMessage(_message,
              ParticipantMessageMonitor.ProcessedMessageState.COMPLETED);
        } else // logging for errors
        {
          code = ErrorCode.ERROR;
          String errorMsg =
              "Message execution failed. msgId: " + getTaskId() + ", errorMsg: "
                  + taskResult.getMessage();
          logger.error(errorMsg);
          _statusUpdateUtil.logError(_message, _handler.getClass(), errorMsg, accessor);
        }
        _executor.getParticipantMonitor().reportProcessedMessage(_message, ParticipantMessageMonitor.ProcessedMessageState.FAILED);
      }

      if (_message.getAttribute(Attributes.PARENT_MSG_ID) == null) {
        // System.err.println("\t[dbg]remove msg: " + getTaskId());
        removeMessageFromZk(accessor, _message);
        reportMessageStat(_manager, _message, taskResult);
        sendReply(accessor, _message, taskResult);
        _executor.finishTask(this);
      }
    } catch (Exception e) {
      exception = e;
      type = ErrorType.FRAMEWORK;
      code = ErrorCode.ERROR;

      String errorMessage =
          "Exception after executing a message, msgId: " + _message.getMsgId() + e;
      logger.error(errorMessage, e);
      _statusUpdateUtil.logError(_message, HelixTask.class, errorMessage, accessor);
      _executor.getParticipantMonitor().reportProcessedMessage(_message, ParticipantMessageMonitor.ProcessedMessageState.FAILED);
    } finally {
      long end = System.currentTimeMillis();
      logger.info("msg: " + _message.getMsgId() + " handling task completed, results:"
          + taskResult.isSuccess() + ", at: " + end + ", took:" + (end - start));

      // Notify the handler about any error happened in the handling procedure, so that
      // the handler have chance to finally cleanup
      if (type == ErrorType.INTERNAL) {
        _handler.onError(taskResult.getException(), code, type);
      } else if (type == ErrorType.FRAMEWORK) {
        _handler.onError(exception, code, type);
      }
    }

    return taskResult;
  }

  private void removeMessageFromZk(HelixDataAccessor accessor, Message message) {
    Builder keyBuilder = accessor.keyBuilder();
    if (message.getTgtName().equalsIgnoreCase("controller")) {
      // TODO: removeProperty returns boolean
      accessor.removeProperty(keyBuilder.controllerMessage(message.getMsgId()));
    } else {
      accessor.removeProperty(keyBuilder.message(_manager.getInstanceName(), message.getMsgId()));
    }
  }

  private void sendReply(HelixDataAccessor accessor, Message message, HelixTaskResult taskResult) {
    if (_message.getCorrelationId() != null
        && !message.getMsgType().equals(MessageType.TASK_REPLY.name())) {
      logger.info("Sending reply for message " + message.getCorrelationId());
      _statusUpdateUtil.logInfo(message, HelixTask.class, "Sending reply", accessor);

      taskResult.getTaskResultMap().put("SUCCESS", "" + taskResult.isSuccess());
      taskResult.getTaskResultMap().put("INTERRUPTED", "" + taskResult.isInterrupted());
      if (!taskResult.isSuccess()) {
        taskResult.getTaskResultMap().put("ERRORINFO", taskResult.getMessage());
      }
      Message replyMessage =
          Message.createReplyMessage(_message, _manager.getInstanceName(),
              taskResult.getTaskResultMap());
      replyMessage.setSrcInstanceType(_manager.getInstanceType());

      if (message.getSrcInstanceType() == InstanceType.PARTICIPANT) {
        Builder keyBuilder = accessor.keyBuilder();
        accessor.setProperty(keyBuilder.message(message.getMsgSrc(), replyMessage.getMsgId()),
            replyMessage);
      } else if (message.getSrcInstanceType() == InstanceType.CONTROLLER) {
        Builder keyBuilder = accessor.keyBuilder();
        accessor.setProperty(keyBuilder.controllerMessage(replyMessage.getMsgId()), replyMessage);
      }
      _statusUpdateUtil.logInfo(message, HelixTask.class,
          "1 msg replied to " + replyMessage.getTgtName(), accessor);
    }
  }

  private void reportMessageStat(HelixManager manager, Message message, HelixTaskResult taskResult) {
    // report stat
    if (!message.getMsgType().equals(MessageType.STATE_TRANSITION.name())) {
      return;
    }
    long now = new Date().getTime();
    long msgReadTime = message.getReadTimeStamp();
    long msgExecutionStartTime = message.getExecuteStartTimeStamp();
    if (msgReadTime != 0 && msgExecutionStartTime != 0) {
      long totalDelay = now - msgReadTime;
      long executionDelay = now - msgExecutionStartTime;
      if (totalDelay > 0 && executionDelay > 0) {
        String fromState = message.getFromState();
        String toState = message.getToState();
        String transition = fromState + "--" + toState;

        StateTransitionContext cxt =
            new StateTransitionContext(manager.getClusterName(), manager.getInstanceName(),
                message.getResourceName(), transition);

        StateTransitionDataPoint data =
            new StateTransitionDataPoint(totalDelay, executionDelay, taskResult.isSuccess());
        _executor.getParticipantMonitor().reportTransitionStat(cxt, data);
      }
    } else {
      logger.warn("message read time and start execution time not recorded.");
    }
  }

  @Override
  public String getTaskId() {
    return _message.getId();
  }

  @Override
  public Message getMessage() {
    return _message;
  }

  @Override
  public NotificationContext getNotificationContext() {
    return _notificationContext;
  }

  @Override
  public void onTimeout() {
    _isTimeout = true;
    _handler.onTimeout();
  }

  @Override
  public synchronized boolean cancel() {
    if (!_isStarted) {
      _isCancelled = true;
      _handler.cancel();
      return true;
    }
    return false;
  }

  private synchronized void setStarted() {
    if (_isCancelled) {
      throw new HelixRollbackException("Task has already been cancelled");
    }
    _isStarted = true;
  }
};
