package org.apache.mesos.hdfs.state;

import com.google.common.base.Predicate;
import com.google.common.collect.Sets;
import com.google.inject.Singleton;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.mesos.Protos;
import org.apache.mesos.hdfs.config.SchedulerConf;
import org.apache.mesos.hdfs.util.HDFSConstants;


import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedHashMap;
import java.util.Set;

@Singleton
public class LiveState {
  public static final Log log = LogFactory.getLog(LiveState.class);
  private Set<Protos.TaskID> stagingTasks = new HashSet<>();
  private AcquisitionPhase currentAcquisitionPhase = AcquisitionPhase.RECONCILING_TASKS;
  // TODO (nicgrayson) Might need to split this out to jns, nns, and dns if dns too big
  private LinkedHashMap<Protos.TaskID, Protos.TaskStatus> runningTasks = new LinkedHashMap<>();
  private HashMap<Protos.TaskStatus, Boolean> nameNode1TaskMap = new HashMap<>();
  private HashMap<Protos.TaskStatus, Boolean> nameNode2TaskMap = new HashMap<>();

  public boolean isNameNode1Initialized() {
    return !nameNode1TaskMap.isEmpty() && nameNode1TaskMap.values().iterator().next();
  }

  public boolean isNameNode2Initialized() {
    return !nameNode2TaskMap.isEmpty() && nameNode2TaskMap.values().iterator().next();
  }

  public void addStagingTask(Protos.TaskID taskId) {
    stagingTasks.add(taskId);
  }

  public int getStagingTasksSize() {
    return stagingTasks.size();
  }

  public void removeStagingTask(final Protos.TaskID taskID) {
    stagingTasks.remove(taskID);
  }

  public LinkedHashMap<Protos.TaskID, Protos.TaskStatus> getRunningTasks() {
    return runningTasks;
  }

  public void removeRunningTask(Protos.TaskID taskId) {
    if (isNameNode1Initialized()
        && nameNode1TaskMap.keySet().iterator().next().getTaskId().equals(taskId)) {
      nameNode1TaskMap.clear();
    } else if (isNameNode2Initialized()
       && nameNode2TaskMap.keySet().iterator().next().getTaskId().equals(taskId)) {
      nameNode2TaskMap.clear();
    }
    runningTasks.remove(taskId);
  }

  public void updateTaskForStatus(Protos.TaskStatus status) {
    //Case of name node, update the task map
    if (status.getTaskId().getValue().contains(HDFSConstants.NAME_NODE_TASKID)) {
      //If initializing the first NN or reconciling the first NN or bootstrapping the first NN
      //set the status to initialized
      if (status.getMessage().equals(HDFSConstants.NAME_NODE_INIT_MESSAGE)
          || (currentAcquisitionPhase.equals(AcquisitionPhase.RECONCILING_TASKS)
              && !isNameNode1Initialized())
          || (status.getMessage().equals(HDFSConstants.NAME_NODE_BOOTSTRAP_MESSAGE)
              && !isNameNode1Initialized())) {
        nameNode1TaskMap.clear();
        nameNode1TaskMap.put(status, true);
        //If bootstrapping the second NN or reconciling the second NN, set the status to initialized
      } else if ((status.getMessage().equals(HDFSConstants.NAME_NODE_BOOTSTRAP_MESSAGE)
          && !isNameNode2Initialized())
          || (currentAcquisitionPhase.equals(AcquisitionPhase.RECONCILING_TASKS)
              && !isNameNode2Initialized())) {
        nameNode2TaskMap.clear();
        nameNode2TaskMap.put(status, true);
        //If the first NN is not running, set the status to running
      } else if (nameNode1TaskMap.isEmpty()) {
        nameNode1TaskMap.put(status, false);
        //If the second NN is not running, set the status to running
      } else if (nameNode2TaskMap.isEmpty()) {
        nameNode2TaskMap.put(status, false);
      }
    }
    runningTasks.put(status.getTaskId(), status);
  }

  public AcquisitionPhase getCurrentAcquisitionPhase() {
    return currentAcquisitionPhase;
  }

  public void transitionTo(AcquisitionPhase phase) {
    this.currentAcquisitionPhase = phase;
  }

  public int getJournalNodeSize() {
    return countOfRunningTasksWith(HDFSConstants.JOURNAL_NODE_ID);
  }

  public int getNameNodeSize() {
    return countOfRunningTasksWith(HDFSConstants.NAME_NODE_TASKID);
  }

  public Protos.TaskID getFirstNameNodeTaskId() {
    if (nameNode1TaskMap.isEmpty()) {
      return null;
    }
    return nameNode1TaskMap.keySet().iterator().next().getTaskId();
  }

  public Protos.TaskID getSecondNameNodeTaskId() {
    if (nameNode2TaskMap.isEmpty()) {
      return null;
    }
    return nameNode2TaskMap.keySet().iterator().next().getTaskId();
  }

  public Protos.SlaveID getFirstNameNodeSlaveId() {
    if (nameNode1TaskMap.isEmpty()) {
      return null;
    }
    return nameNode1TaskMap.keySet().iterator().next().getSlaveId();
  }

  public Protos.SlaveID getSecondNameNodeSlaveId() {
    if (nameNode2TaskMap.isEmpty()) {
      return null;
    }
    return nameNode2TaskMap.keySet().iterator().next().getSlaveId();
  }

  private int countOfRunningTasksWith(final String nodeId) {
    return Sets.filter(runningTasks.keySet(), new Predicate<Protos.TaskID>() {
      @Override
      public boolean apply(Protos.TaskID taskID) {
        return taskID.getValue().contains(nodeId);
      }
    }).size();
  }
}

@@ -2219,16 +2219,356 @@
 %7B%0A    //
+ TODO (elingg) Use Starting Status when the task is running, but not initialized. Use running%0A    // status when the task is initialized so that we can differentiate during the reconciliation%0A    // phase. Also, add the health checks which will kill the task if it doesn't properly%0A    // initialize or if it reaches an error state.%0A    // 
 Case of 
@@ -2688,16 +2688,17 @@
       //
+ 
 If initi
@@ -2784,16 +2784,17 @@
       //
+ 
 set the 
@@ -3230,20 +3230,21 @@
 ;%0A      
- 
+%7D
  //
+ 
 If boots
@@ -3295,16 +3295,27 @@
 cond NN,
+%0A        //
  set the
@@ -3338,26 +3338,24 @@
 alized%0A     
- %7D
  else if ((s
@@ -3671,28 +3671,29 @@
 rue);%0A      
- 
+%7D
  //
+ 
 If the first
@@ -3736,33 +3736,32 @@
 o running%0A      
-%7D
  else if (nameNo
@@ -3835,20 +3835,21 @@
 ;%0A      
- 
+%7D
  //
+ 
 If the s
@@ -3901,25 +3901,24 @@
 nning%0A      
-%7D
  else if (na

Add TODO/ Update javadoc