package com.neverwinterdp.vm.environment.yarn;

import java.net.InetAddress;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.atomic.AtomicInteger;

import org.apache.hadoop.yarn.api.protocolrecords.AllocateResponse;
import org.apache.hadoop.yarn.api.protocolrecords.RegisterApplicationMasterResponse;
import org.apache.hadoop.yarn.api.records.Container;
import org.apache.hadoop.yarn.api.records.ContainerExitStatus;
import org.apache.hadoop.yarn.api.records.ContainerState;
import org.apache.hadoop.yarn.api.records.ContainerStatus;
import org.apache.hadoop.yarn.api.records.FinalApplicationStatus;
import org.apache.hadoop.yarn.api.records.NodeReport;
import org.apache.hadoop.yarn.client.api.AMRMClient;
import org.apache.hadoop.yarn.client.api.async.AMRMClientAsync;

import com.google.inject.Inject;
import com.google.inject.Singleton;
import com.mycila.jmx.annotation.JmxBean;
import com.neverwinterdp.vm.VMConfig;

@Singleton
@JmxBean("role=vm-manager, type=YarnManager, name=AsynYarnManager")
public class SyncYarnManager extends YarnManager {
  private AMRMClient<ContainerRequest> amrmClient ;
  
  private AtomicInteger countContainerRequest = new AtomicInteger();
  
  
  public AMRMClient<ContainerRequest> getAMRMClient() { return this.amrmClient ; }
  
  
  @Inject
  public void onInit(VMConfig vmConfig) throws Exception {
    logger.info("Start init(VMConfig vmConfig)");
    try {
      super.onInit(vmConfig);
      amrmClient = AMRMClient.createAMRMClient();
      amrmClient.init(conf);
      amrmClient.start();
      // Register with RM
      String appHostName = InetAddress.getLocalHost().getHostAddress()  ;
      RegisterApplicationMasterResponse registerResponse = amrmClient.registerApplicationMaster(appHostName, 0, "");
    } catch(Throwable t) {
      logger.error("Error: " , t);
      t.printStackTrace();
      throw t;
    }
    logger.info("Finish init(VMConfig vmConfig)");
  }

  public void onDestroy()  {
    logger.info("Start onDestroy()");
    try {
      if(amrmClient != null) {
        amrmClient.unregisterApplicationMaster(FinalApplicationStatus.SUCCEEDED, "", "");
        amrmClient.stop();
        amrmClient.close(); 
      }
      super.onDestroy();
    } catch(Exception ex) {
      logger.error("Cannot destroy YarnManager properly", ex);
    }
    logger.info("Finish onDestroy()");
  }
  
  synchronized public void add(ContainerRequest containerReq, ContainerRequestCallback callback) throws Exception {
    logger.info("Start add count = " + countContainerRequest.incrementAndGet());
    containerReq.setCallback(callback);
    amrmClient.addContainerRequest(containerReq);
    int allocatedContainers = 0;
    long stopTime = System.currentTimeMillis() + 180000;
    int retry = 0;
    while (allocatedContainers < 1 && System.currentTimeMillis() < stopTime) {
      retry++ ;
      AllocateResponse response = amrmClient.allocate(0 /*progress indicator*/);
      List<Container> containers = response.getAllocatedContainers();
      logger.info("  " + retry + "  Allocate containers = " + containers.size() + ", retry = " + retry + ", duration = " + (retry * 500));
      for (Container container : containers) {
        callback.onAllocate(this, containerReq, container);
        allocatedContainers++;
        break;
      }
      Thread.sleep(500);
    }
    amrmClient.removeContainerRequest(containerReq);
    if(allocatedContainers == 0) {
      throw new Exception("Cannot allocate the container");
    }
    logger.info("Finish add");
  }
}
@@ -74,36 +74,8 @@
 ss;%0A
-import java.util.ArrayList;%0A
 impo
@@ -396,381 +396,83 @@
 rds.
-ContainerExitStatus;%0Aimport org.apache.hadoop.yarn.api.records.ContainerState;%0Aimport org.apache.hadoop.yarn.api.records.ContainerStatus;%0Aimport org.apache.hadoop.yarn.api.records.FinalApplicationStatus;%0Aimport org.apache.hadoop.yarn.api.records.NodeReport;%0Aimport org.apache.hadoop.yarn.client.api.AMRMClient;%0Aimport org.apache.hadoop.yarn.client.api.async.AMRMClientAsync
+FinalApplicationStatus;%0Aimport org.apache.hadoop.yarn.client.api.AMRMClient
 ;%0A%0Ai
@@ -2283,24 +2283,493 @@
 lback);%0A    
+Container container = null ;%0A    for(int i = 0; i %3C 5; i++) %7B%0A      container = allocate(containerReq, 5000);%0A      if(container != null) %7B%0A        break;%0A      %7D%0A    %7D%0A    if(container != null) %7B%0A      callback.onAllocate(this, containerReq, container);%0A    %7D else %7B%0A      throw new Exception(%22Cannot allocate the container%22);%0A    %7D%0A    logger.info(%22Finish add%22);%0A  %7D%0A  %0A  private Container allocate(ContainerRequest containerReq, long timeout) throws Exception %7B%0A    
 amrmClient.a
@@ -2804,41 +2804,8 @@
 Req)
-;%0A    int allocatedContainers = 0
 ;%0A  
@@ -2855,14 +2855,15 @@
 ) + 
-180000
+timeout
 ;%0A  
@@ -2894,35 +2894,8 @@
 le (
-allocatedContainers %3C 1 && 
 Syst
@@ -3296,104 +3296,24 @@
     
-callback.onAllocate(this, containerReq, container);%0A        allocatedContainers++;%0A        break
+return container
 ;%0A  
@@ -3410,134 +3410,20 @@
     
-if(allocatedContainers == 0) %7B%0A      throw new Exception(%22Cannot allocate the container%22);%0A    %7D%0A    logger.info(%22Finish add%22)
+return null 
 ;%0A  

add retry for container request