/**
 * 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.ambari.server.state.scheduler;

import java.util.HashMap;
import java.util.Map;

import org.apache.ambari.server.AmbariException;
import org.apache.ambari.server.actionmanager.HostRoleStatus;
import org.apache.ambari.server.scheduler.AbstractLinearExecutionJob;
import org.apache.ambari.server.scheduler.ExecutionScheduleManager;
import org.quartz.DisallowConcurrentExecution;
import org.quartz.PersistJobDataAfterExecution;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.inject.Inject;
import com.google.inject.name.Named;

@PersistJobDataAfterExecution
@DisallowConcurrentExecution
public class BatchRequestJob extends AbstractLinearExecutionJob {
  private static final Logger LOG = LoggerFactory.getLogger(BatchRequestJob.class);

  public static final String BATCH_REQUEST_EXECUTION_ID_KEY =
    "BatchRequestJob.ExecutionId";
  public static final String BATCH_REQUEST_BATCH_ID_KEY =
    "BatchRequestJob.BatchId";
  public static final String BATCH_REQUEST_CLUSTER_NAME_KEY =
    "BatchRequestJob.ClusterName";
  public static final String BATCH_REQUEST_FAILED_TASKS_KEY =
    "BatchRequestJob.FailedTaskCount";
  public static final String BATCH_REQUEST_TOTAL_TASKS_KEY =
    "BatchRequestJob.TotalTaskCount";

  private final long statusCheckInterval;

  @Inject
  public BatchRequestJob(ExecutionScheduleManager executionScheduleManager,
                         @Named("statusCheckInterval") long statusCheckInterval) {
    super(executionScheduleManager);
    this.statusCheckInterval = statusCheckInterval;
  }

  @Override
  protected void doWork(Map<String, Object> properties) throws AmbariException {

    Long executionId = properties.get(BATCH_REQUEST_EXECUTION_ID_KEY) != null ?
      (Long) properties.get(BATCH_REQUEST_EXECUTION_ID_KEY) : null;
    Long batchId = properties.get(BATCH_REQUEST_BATCH_ID_KEY) != null ?
      (Long) properties.get(BATCH_REQUEST_BATCH_ID_KEY) : null;
    String clusterName = (String) properties.get(BATCH_REQUEST_CLUSTER_NAME_KEY);


    if (executionId == null || batchId == null) {
      throw new AmbariException("Unable to retrieve persisted batch request"
        + ", execution_id = " + executionId
        + ", batch_id = " + batchId);
    }

    // Aggregate tasks counts stored in the DataMap
    Map<String, Integer> taskCounts = getTaskCountProperties(properties);

    Long requestId = executionScheduleManager.executeBatchRequest
      (executionId, batchId, clusterName);

    if (requestId != null) {
      HostRoleStatus status;
      BatchRequestResponse batchRequestResponse;
      do {
        batchRequestResponse = executionScheduleManager
          .getBatchRequestResponse(requestId, clusterName);

        status = HostRoleStatus.valueOf(batchRequestResponse.getStatus());

        executionScheduleManager.updateBatchRequest(executionId, batchId,
          clusterName, batchRequestResponse, true);

        try {
          Thread.sleep(statusCheckInterval);
        } catch (InterruptedException e) {
          String message = "Job Thread interrupted";
          LOG.error(message, e);
          throw new AmbariException(message, e);
        }
      } while (!status.isCompletedState());

      // Store aggregated task status counts in the DataMap
      Map<String, Integer> aggregateCounts = addTaskCountToProperties
        (properties, taskCounts, batchRequestResponse);

      if (executionScheduleManager.hasToleranceThresholdExceeded
          (executionId, clusterName, aggregateCounts)) {

        throw new AmbariException("Task failure tolerance limit exceeded"
            + ", execution_id = " + executionId
            + ", processed batch_id = " + batchId
            + ", failed tasks = " + aggregateCounts.get(BATCH_REQUEST_FAILED_TASKS_KEY)
            + ", total tasks completed = " + aggregateCounts.get(BATCH_REQUEST_TOTAL_TASKS_KEY));
      }
    }
  }

  @Override
  protected void finalizeExecution(Map<String, Object> properties)
      throws AmbariException {

    Long executionId = properties.get(BATCH_REQUEST_EXECUTION_ID_KEY) != null ?
      (Long) properties.get(BATCH_REQUEST_EXECUTION_ID_KEY) : null;
    Long batchId = properties.get(BATCH_REQUEST_BATCH_ID_KEY) != null ?
      (Long) properties.get(BATCH_REQUEST_BATCH_ID_KEY) : null;
    String clusterName = (String) properties.get(BATCH_REQUEST_CLUSTER_NAME_KEY);

    if (executionId == null || batchId == null) {
      throw new AmbariException("Unable to retrieve persisted batch request"
        + ", execution_id = " + executionId
        + ", batch_id = " + batchId);
    }

    // Check if this job has a future and update status if it doesn't
    executionScheduleManager.finalizeBatch(executionId, clusterName);

  }

  private Map<String, Integer> addTaskCountToProperties(Map<String, Object> properties,
                                        Map<String, Integer> oldCounts,
                                        BatchRequestResponse batchRequestResponse) {

    Map<String, Integer> taskCounts = new HashMap<>();

    if (batchRequestResponse != null) {
      Integer failedTasks = batchRequestResponse.getFailedTaskCount() +
        batchRequestResponse.getAbortedTaskCount() +
        batchRequestResponse.getTimedOutTaskCount();

      Integer failedCount = oldCounts.get(BATCH_REQUEST_FAILED_TASKS_KEY) + failedTasks;
      Integer totalCount = oldCounts.get(BATCH_REQUEST_TOTAL_TASKS_KEY) +
        batchRequestResponse.getTotalTaskCount();

      taskCounts.put(BATCH_REQUEST_FAILED_TASKS_KEY, failedCount);
      taskCounts.put(BATCH_REQUEST_TOTAL_TASKS_KEY, totalCount);

      properties.put(BATCH_REQUEST_FAILED_TASKS_KEY, failedCount);
      properties.put(BATCH_REQUEST_TOTAL_TASKS_KEY, totalCount);
    }

    return taskCounts;
  }

  private Map<String, Integer> getTaskCountProperties(Map<String, Object> properties) {
    Map<String, Integer> taskCounts = new HashMap<>();
    if (properties != null) {
      Object countObj = properties.get(BATCH_REQUEST_FAILED_TASKS_KEY);
      taskCounts.put(BATCH_REQUEST_FAILED_TASKS_KEY,
        countObj != null ? Integer.parseInt(countObj.toString()) : 0);
      countObj = properties.get(BATCH_REQUEST_TOTAL_TASKS_KEY);
      taskCounts.put(BATCH_REQUEST_TOTAL_TASKS_KEY, countObj != null ?
        Integer.parseInt(countObj.toString()) : 0);
    }
    return taskCounts;
  }
}