/**
 *    Copyright (C) 2015 Mesosphere, Inc.
 *
 * Licensed 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 io.mesosphere.mesos.frameworks.cassandra.scheduler;

import com.google.common.base.Optional;
import io.mesosphere.mesos.frameworks.cassandra.CassandraFrameworkProtos;
import org.jetbrains.annotations.NotNull;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.List;

public class NodeTaskClusterJobHandler extends ClusterJobHandler {
    private static final Logger LOGGER = LoggerFactory.getLogger(NodeTaskClusterJobHandler.class);

    public NodeTaskClusterJobHandler(@NotNull final CassandraCluster cluster, @NotNull final PersistedCassandraClusterJobs jobsState) {
        super(cluster, jobsState);
    }

    @Override
    public void handleTaskOffer(@NotNull final CassandraFrameworkProtos.ClusterJobStatus currentJob, @NotNull final String executorId, @NotNull final Optional<CassandraFrameworkProtos.CassandraNode> nodeForExecutorId, @NotNull final TasksForOffer tasksForOffer) {
        if (currentJob.hasCurrentNode()) {
            final CassandraFrameworkProtos.NodeJobStatus nodJobStatus = currentJob.getCurrentNode();
            if (executorId.equals(nodJobStatus.getExecutorId())) {
                // submit status request
                tasksForOffer.getSubmitTasks().add(CassandraFrameworkProtos.TaskDetails.newBuilder()
                    .setType(CassandraFrameworkProtos.TaskDetails.TaskDetailsType.NODE_JOB_STATUS)
                    .build());

                LOGGER.info("Inquiring cluster job status for {} from {}", currentJob.getJobType().name(),
                    nodJobStatus.getExecutorId());

                return;
            }
            return;
        }

        if (currentJob.getAborted() && !currentJob.hasCurrentNode()) {
            jobsState.setCurrentJob(null);
            // TODO record aborted job in history??
            return;
        }

        if (!currentJob.hasCurrentNode()) {
            final List<String> remainingNodes = new ArrayList<>(currentJob.getRemainingNodesList());
            if (remainingNodes.isEmpty()) {
                jobsState.finishJob(currentJob);
                return;
            }

            if (!remainingNodes.remove(executorId)) {
                return;
            }

            if (!nodeForExecutorId.isPresent()) {
                rejectNode(remainingNodes);
                return;
            }

            final CassandraFrameworkProtos.CassandraNode node = nodeForExecutorId.get();

            if (node.getTargetRunState() != CassandraFrameworkProtos.CassandraNode.TargetRunState.RUN
                || !cluster.isLiveNode(node)) {
                rejectNode(remainingNodes);
                return;
            }

            final CassandraFrameworkProtos.TaskDetails taskDetails = CassandraFrameworkProtos.TaskDetails.newBuilder()
                .setType(CassandraFrameworkProtos.TaskDetails.TaskDetailsType.NODE_JOB)
                .setNodeJobTask(CassandraFrameworkProtos.NodeJobTask.newBuilder().setJobType(currentJob.getJobType()))
                .build();
            final CassandraFrameworkProtos.CassandraNodeTask cassandraNodeTask = CassandraFrameworkProtos.CassandraNodeTask.newBuilder()
                .setType(CassandraFrameworkProtos.CassandraNodeTask.NodeTaskType.CLUSTER_JOB)
                .setTaskId(executorId + '.' + currentJob.getJobType().name())
                .setResources(CassandraFrameworkProtos.TaskResources.newBuilder()
                    .setCpuCores(0.1)
                    .setMemMb(16)
                    .setDiskMb(16))
                .setTaskDetails(taskDetails)
                .build();
            tasksForOffer.getLaunchTasks().add(cassandraNodeTask);

            final CassandraFrameworkProtos.NodeJobStatus currentNode = CassandraFrameworkProtos.NodeJobStatus.newBuilder()
                .setExecutorId(node.getCassandraNodeExecutor().getExecutorId())
                .setTaskId(cassandraNodeTask.getTaskId())
                .setJobType(currentJob.getJobType())
                .setStartedTimestamp(System.currentTimeMillis())
                .build();
            jobsState.updateJobCurrentNode(currentJob, currentNode);

            LOGGER.info("Starting cluster job {} on {}/{}", currentJob.getJobType().name(), node.getIp(),
                node.getHostname());
        }
    }

    private void rejectNode(@NotNull final List<String> remainingNodes) {
        final CassandraFrameworkProtos.ClusterJobStatus currentJob;
        currentJob = CassandraFrameworkProtos.ClusterJobStatus.newBuilder()
            .clearRemainingNodes()
            .addAllRemainingNodes(remainingNodes)
            .build();
        jobsState.setCurrentJob(currentJob);
    }

    @Override
    public void onNodeJobStatus(@NotNull final CassandraFrameworkProtos.ClusterJobStatus currentJob, @NotNull final CassandraFrameworkProtos.NodeJobStatus nodeJobStatus) {
        LOGGER.info("Got node job status from {}, running={}", nodeJobStatus.getExecutorId(), nodeJobStatus.getRunning());

        if (currentJob.getCurrentNode() != null && currentJob.getCurrentNode().getExecutorId().equals(nodeJobStatus.getExecutorId())) {
            if (nodeJobStatus.getRunning()) {
                jobsState.setCurrentJob(CassandraFrameworkProtos.ClusterJobStatus.newBuilder(currentJob)
                    .setCurrentNode(nodeJobStatus)
                    .build());
            } else {
                nodeFinished(nodeJobStatus, currentJob);
            }
        }
    }
}

@@ -2898,32 +2898,44 @@
      rejectNode(
+currentJob, 
 remainingNodes);
@@ -3241,16 +3241,28 @@
 ectNode(
+currentJob, 
 remainin
@@ -4973,24 +4973,93 @@
 tNull final 
+CassandraFrameworkProtos.ClusterJobStatus currentJob, @NotNull final 
 List%3CString%3E
@@ -5137,40 +5137,36 @@
 tus 
-currentJob;%0A        currentJob =
+updatedJob =%0A               
  Cas
@@ -5207,34 +5207,52 @@
 atus.newBuilder(
-)%0A
+currentJob)%0A        
             .cle
@@ -5247,16 +5247,20 @@
         
+    
 .clearRe
@@ -5274,16 +5274,28 @@
 Nodes()%0A
+            
         
@@ -5340,32 +5340,44 @@
 es)%0A            
+   
+         
 .build();%0A      
@@ -5398,31 +5398,31 @@
 tCurrentJob(
-current
+updated
 Job);%0A    %7D%0A

Update code to ensure that currently running job is used a base of building new job state.