public class CommandLineHandler  {
    private static final Logger LOGGER = LoggerFactory.getLogger(CommandLineHandler.class);

    private static final String OPTION_HELP = "h";
    private static final String OPTION_VERSION = "v";
    private static final String OPTION_PORT = "P";
    private static final String OPTION_BIND = "b";
    private static final String OPTION_MASTER = "m";
    private static final String OPTION_NAME = "n";
    private static final String OPTION_ZOOKEEPERS = "z";
    private static final String OPTION_TARBALL = "t";
    private static final String OPTION_FRAMEWORK = "fc";
    private static final String OPTION_CLUSTER = "cc";
    private static final String OPTION_INITIALIZE = "i";

    private static final Options options;
    static {
        options = new Options();
        options.addOption(OPTION_HELP, "help", false, "Print this message and exit");
        options.addOption(OPTION_VERSION, "version", false, "Show version number");
        options.addOption(OPTION_PORT, "port", true, "Port number to serve HTTP interface");
        options.addOption(OPTION_BIND, "bind-address", true, "IP address of interface to bind HTTP interface to");
        options.addOption(OPTION_MASTER, "master", true, "Location of mesos master to connect to");
        options.addOption(OPTION_NAME, "name", true, "Name of this mesos framework");
        options.addOption(OPTION_ZOOKEEPERS, "zookeepers", true, "List of Zookeeper servers");
        options.addOption(OPTION_TARBALL, "tarball", true, "URI of framework tarball");
        options.addOption(OPTION_FRAMEWORK, "framework", true, "JSON file of entire framework configuration");
        options.addOption(OPTION_CLUSTER, "cluster", true, "JSON file containing cluster configuration");
        options.addOption(OPTION_INITIALIZE, "init", false, "If present, initialize new Accumulo instance");
    }

    private static final Parser parser = new GnuParser();
    private final CommandLine cmdLine;

    public CommandLineHandler(CommandLine cmdLine){
        this.cmdLine = cmdLine;
    }

    public CommandLineHandler(String args[]){
        this(CommandLineHandler.parseArgs(args));
    }

    /**
     * Parses command line arguments for configuration information.
     *
     * Explicit options override anything in JSON. For example if JSON specifies
     * port as 8080 and commandline has -P 5150, the port will be 5150. The returned
     * object might not contain enough information to initialize and run a cluster.
     * This is because a recovery re-launch of the framework should just be able to
     * read the configuration from the Mesos store.
     *
     * @return Framework configuration
     *
     */
    public Framework getFrameworkDefinition() throws IOException {

        ObjectMapper mapper = new ObjectMapper();

        Framework config = null;
        // Read full config if available
        if( cmdLine.hasOption(OPTION_FRAMEWORK)){
            File frameworkFile = new File(cmdLine.getOptionValue(OPTION_FRAMEWORK));
            LOGGER.info("Reading Framework Configuration from: {}", frameworkFile.getAbsolutePath());
            config = mapper.readValue(frameworkFile, Framework.class);
        } else {
            config = new Framework();
        }

        // Override config with explicit command line options.
        if( cmdLine.hasOption(OPTION_PORT) ){
            config.setHttpPort(Integer.parseInt(cmdLine.getOptionValue(OPTION_PORT)));
        }

        if( cmdLine.hasOption(OPTION_BIND) ){
            config.setBindAddress(cmdLine.getOptionValue(OPTION_BIND));
        }

        if ( cmdLine.hasOption(OPTION_MASTER) ){
            config.setMesosMaster(cmdLine.getOptionValue(OPTION_MASTER));
        }

        if ( cmdLine.hasOption(OPTION_NAME) ){
            config.setName(cmdLine.getOptionValue(OPTION_NAME));
        }

        if (cmdLine.hasOption(OPTION_ZOOKEEPERS) ){
            config.setZkServers(cmdLine.getOptionValue(OPTION_ZOOKEEPERS));
        }

        if (cmdLine.hasOption(OPTION_TARBALL)){
            config.setTarballUri(cmdLine.getOptionValue(OPTION_TARBALL));
        }

        // If there's an Accumulo cluster JSON provided, override.
        if (cmdLine.hasOption(OPTION_CLUSTER)) {
            File clusterFile = new File(cmdLine.getOptionValue(OPTION_CLUSTER));
            LOGGER.info("Loading cluster configuration {}", clusterFile.getAbsolutePath());
            Accumulo clusterConfig = mapper.readValue(clusterFile, Accumulo.class);
            config.setCluster(clusterConfig);
        }

        LOGGER.debug("Configuration after command line handling: " + config.toString());

        return config;
    }

    public boolean isInitializeRequest(){ return cmdLine.hasOption(OPTION_INITIALIZE); }

    private static CommandLine parseArgs(String args[]) {
        CommandLine cmdLine = null;
        try {
            cmdLine = parser.parse(options, args);
        } catch (ParseException e) {
            System.err.println(e);
        }
        return cmdLine;
    }

    // print help or version information.
    public boolean checkHelpOrVersion() {

        // Handle exiting options
        if (cmdLine.hasOption('h')){
            printHelp();
            return true;
        }
        if (cmdLine.hasOption('v')){
            printVersion();
            return true;
        }
        return false;
    }

    private void printVersion(){
        System.out.format("%s : %s\n", Constants.EXE_NAME, Constants.FRAMEWORK_VERSION);
    }

    private void printHelp(){
        HelpFormatter formatter = new HelpFormatter();
        formatter.printHelp( Constants.EXE_NAME, options, true );
    }

} public final class Constants {

    // bundler used for items defined by Maven
    private static final ResourceBundle rb = ResourceBundle.getBundle("accumulo-mesos");
    public static final String FRAMEWORK_VERSION = rb.getString("application.version");
    public static final String EXE_NAME = rb.getString("application.name");
    public static final String EXECUTOR_JAR = rb.getString("application.name.executor")
            +"-"
            +FRAMEWORK_VERSION
            +"-jar-with-dependencies"
            +".jar";

    public static final String ACCUMULO_MESOS_DISTRO = "accumulo-mesos-dist-"+FRAMEWORK_VERSION;
    public static final String ACCUMULO_NATIVE_LIB = "libaccumulo.so";
} public final class Defaults {

    private static final ResourceBundle rb = ResourceBundle.getBundle("accumulo-mesos");

    // Default Port to bind http interface to
    public final static int HTTP_PORT = 18120;

    // Bind to all interfaces by default
    public final static String BIND_ADDRESS = "0.0.0.0";

    // Look for mesos master on local machine
    public final static String MESOS_MASTER = "10.0.2.15:5050";

    // Framework name
    public final static String FRAMEWORK_NAME = "Accumulo-Mesos";

    // Zookeeper server
    public final static String ZK_SERVERS = "127.0.0.1:2181";
    // Zookeeper connection timeout
    public final static long ZOOKEEPER_TIMEOUT = 3l;
    public final static TimeUnit ZK_TIMEOUT_UNIT = TimeUnit.SECONDS;

    // I spent way too much time figuring out this needs a leading / so don't be me
    public final static String ZK_STATE_ZNODE = "/accumulo-mesos";

    // Memory settings are MB
    public final static int MIN_MASTER_MEM = 512;
    public final static int MIN_TSERVER_MEM = 512;
    public final static int MIN_GC_MEM = 128;
    public final static int MIN_MONITOR_MEM = 128;
    
    public final static String DEFAULT_8G_MAX_MEMORY = "8192";
    public final static String DEFAULT_2G_MAX_MEMORY = "2048";
    public final static String DEFAULT_1G_MAX_MEMORY = "1024";

    // Minimum CPU settings
    public final static int MIN_MASTER_CPUS = 1;
    public final static int MIN_TSERVER_CPUS = 1;
    public final static int MIN_GC_CPUS = 1;
    public final static int MIN_MONITOR_CPUS = 1;

    // Minimum number of tservers to launch
    public final static int MIN_TSERVERS = 3;

    public final static String ROOT_USER = "root";
    public final static String ROOT_PASSWORD = "secret";

    public final static double MAX_EXECUTOR_MEM = 128;
    public final static double MIN_EXECUTOR_MEM = 16;
    // CPUS to be consumed by executor outside of server process
    public final static double EXECUTOR_CPUS = 0.25;
    
    // Default location of the accumulo-site.xml file
    public final static String ACCUMULO_SITE_URI = "file:./conf/default-accumulo-site.xml";
    

} public class Environment {

    public static final String JAVA_HOME = "JAVA_HOME";

    public static final String CLASSPATH = "CLASSPATH";

    public static final String HADOOP_PREFIX = "HADOOP_PREFIX";
    public static final String HADOOP_CONF_DIR = "HADOOP_CONF_DIR";
    public static final String HADOOP_HOME = "HADOOP_HOME";
    
    public static final String ACCUMULO_HOME = "ACCUMULO_HOME";
    public static final String ACCUMULO_LOG_DIR = "ACCUMULO_LOG_DIR";
    public static final String ACCUMULO_CLIENT_CONF_PATH = "ACCUMULO_CLIENT_CONF_PATH";
    public static final String ACCUMULO_CONF_DIR = "ACCUMULO_CONF_DIR";
    public static final String ACCUMULO_WALOG = "ACCUMULO_WALOG";
   
    public static final String LD_LIBRARY_PATH = "LD_LIBRARY_PATH";
    public static final String DYLD_LIBRARY_PATH = "DYLD_LIBRARY_PATH";
   
    // List of paths separated by commas
    public static final String NATIVE_LIB_PATHS = "NATIVE_LIB_PATHS";
    
    public static final String ZOOKEEPER_HOME = "ZOOKEEPER_HOME";

    public static final String MESOS_DIRECTORY = "MESOS_DIRECTORY";

    public static final List<String> REQUIRED_FRAMEWORK_VARS = Arrays.asList(
                   JAVA_HOME,
                   ACCUMULO_HOME,
                   HADOOP_PREFIX,
                   HADOOP_CONF_DIR,
                   ZOOKEEPER_HOME);

    public static final List<String> REQUIRED_EXECUTOR_VARS = Arrays.asList(
                   JAVA_HOME,
                   ACCUMULO_HOME,
                   HADOOP_PREFIX,
                   ZOOKEEPER_HOME,
                   ACCUMULO_LOG_DIR);


    /*
        Optional Accumulo vars for executor
     */
    public static final String ACCUMULO_XTRAJARS = "ACCUMULO_XTRAJARS";
    public static final String ACCUMULO_GENERAL_OPTS = "ACCUMULO_GENERAL_OPTS";
    public static final String ACCUMULO_JAAS_CONF = "ACCUMULO_JAAS_CONF"; //-Djava.security.auth.login.config=${ACCUMULO_JAAS_CONF}
    public static final String ACCUMULO_KRB5_CONF = "ACCUMULO_KRB5_CONF"; //-Djava.security.krb5.conf=${ACCUMULO_KRB5_CONF}

    public static final String ACCUMULO_MASTER_OPTS = "ACCUMULO_MASTER_OPTS";
    public static final String ACCUMULO_GC_OPTS = "ACCUMULO_GC_OPTS";
    public static final String ACCUMULO_TSERVER_OPTS = "ACCUMULO_TSERVER_OPTS";
    public static final String ACCUMULO_MONITOR_OPTS = "ACCUMULO_MONITOR_OPTS";
    public static final String ACCUMULO_LOGGER_OPTS = "ACCUMULO_LOGGER_OPTS";
    public static final String ACCUMULO_OTHER_OPTS = "ACCUMULO_OTHER_OPTS";
    public static final String ACCUMULO_KILL_CMD = "ACCUMULO_KILL_CMD";

    // generated
    public static final String START_JAR = "START_JAR";
    public static final String LIB_PATH = "LIB_PATH";
    /*
    if [ -e "${HADOOP_PREFIX}/lib/native/libhadoop.so" ]; then
            LIB_PATH="${HADOOP_PREFIX}/lib/native"
    LD_LIBRARY_PATH="${LIB_PATH}:${LD_LIBRARY_PATH}"     # For Linux
    DYLD_LIBRARY_PATH="${LIB_PATH}:${DYLD_LIBRARY_PATH}" # For Mac
    fi
    */


    public static List<String> getMissingVariables(List<String> vars) {
        List<String> missingVariables = new ArrayList<>();
        Set<String> envKeys = System.getenv().keySet();
        for(String var : vars ) {
            if(!envKeys.contains(var)){
                missingVariables.add(var);
            }
        }
        return missingVariables;
    }

    public static String get(String var){
        return System.getenv(var);
    }

} public class WebServer {

    private final Framework config;
    private Server server;

    @Inject
    public WebServer(Framework config) {
        this.config = config;
    }

    public void start() throws java.lang.Exception {

        checkParam(this.config.getBindAddress(), "Bind address");
        checkParam(this.config.getHttpPort() + "", "Http Port");
       
        this.server = new Server();
        // configure server
        ServerConnector http = new ServerConnector(server);
        http.setHost(this.config.getBindAddress());
        http.setPort(this.config.getHttpPort());
        server.addConnector(http);

        ServletContextHandler context = new ServletContextHandler(this.server, "/", ServletContextHandler.NO_SESSIONS);
        context.addFilter(GuiceFilter.class, "/*", EnumSet.of(javax.servlet.DispatcherType.REQUEST, javax.servlet.DispatcherType.ASYNC));
        context.addServlet(DefaultServlet.class, "/*");

        String staticDir = this.getClass().getClassLoader().getResource("webapp/public").toExternalForm();
        context.setResourceBase(staticDir);

        this.server.start();
    }

    public void stop() throws java.lang.Exception{
        this.server.stop();
    }
    
    private void checkParam(String value, String name) {
        if (StringUtils.isEmpty(value)) {
            throw new RuntimeException("Required parameter " + name + " is missing");
        }
    }
} public class ApiServletModule extends ServletModule {

    @Override
    protected void configureServlets() {
        bind(DefaultServlet.class).in(Singleton.class);

        bind(EchoResource.class);
        bind(ClusterApi.class);
        bind(ConfigApi.class);
        bind(DefaultApi.class);
        bind(StatusApi.class);

        serve("/api/*").with(GuiceContainer.class);
    }

} public class ConfigurationModule extends AbstractModule {

    private final Framework config;

    public ConfigurationModule(Framework config){
        this.config = config;
    }

    @Override
    protected void configure() {
        bind(Framework.class).toInstance(config);
    }

} public class AccumuloInitializer {
    private static final Logger LOGGER = LoggerFactory.getLogger(AccumuloInitializer.class);

    private Accumulo config;
    private String accumuloHome;
    
    public AccumuloInitializer(Accumulo config) {
        this.config = config;
        this.accumuloHome = Environment.get(Environment.ACCUMULO_HOME);
    }
    
    /**
     * Write the accumulo site file and initialize accumulo.
     * 
     * The system property Environment.ACCUMULO_HOME and HADOOP_CONF_DIR must be set and config must have
     * the accumulo instance name, root password, along with the accumulo directories set.
     *
     * @return accumulo instance name
     */
    public int initialize() throws IOException {
        
        // run accumulo init procedure
        LOGGER.info("Writing accumulo-site.xml");

        AccumuloSiteXml siteXml = new AccumuloSiteXml(this.config);
        siteXml.initializeFromScheduler(AccumuloSiteXml.getEmptySiteXml());

        writeAccumuloSiteFile(accumuloHome, siteXml);
        config.setSiteXml(siteXml.toString());

        // accumulo-env.sh
        copyAccumuloEnvFile(accumuloHome); // IOException

        LinkedList<String> initArgs  = new LinkedList<>();
        initArgs.add("--instance-name");
        initArgs.add(config.getInstance());
        initArgs.add("--password");
        initArgs.add(config.getRootPassword());
        initArgs.add("--clear-instance-name");

        AccumuloProcessFactory processFactory = new AccumuloProcessFactory();
       
        Process initProcess;
        int status = 0;
        try {
            initProcess = processFactory.exec(ServerProfile.TypeEnum.init.getServerKeyword(),
                                              initArgs.toArray(new String[initArgs.size()]));
            initProcess.waitFor();
            status = initProcess.exitValue();
        } catch (Exception ioe) {
            LOGGER.error("IOException while trying to initialize Accumulo", ioe);
            status = -1;
        }  
        return status;
    }

    public static void writeAccumuloSiteFile(String accumuloHomeDir, AccumuloSiteXml siteXml) {
        LOGGER.info("ACCUMULO HOME? " + accumuloHomeDir);
        try {
  
            File accumuloSiteFile = new File(accumuloHomeDir + File.separator +
                    "conf" + File.separator + "accumulo-site.xml");

            LOGGER.info("Writing accumulo-site.xml to {}", accumuloSiteFile.getAbsolutePath());

            OutputStream siteFile = new FileOutputStream(accumuloSiteFile);
            IOUtils.write(siteXml.toXml(), siteFile);
            IOUtils.closeQuietly(siteFile);

        } catch (Exception e) {
            logErrorAndDie("Error Creating accumulo-site.xml\n",e);
        }
    }

    /**
     * This is required because the bin/accumulo script complains without it. The framework should be setting
     * all the environment variables that this sets, and the script uses them if already set.
     *
     * @param accumuloHomeDir
     * @throws IOException
     */
    public static void copyAccumuloEnvFile(String accumuloHomeDir) throws IOException {
        File inputFile = new File(accumuloHomeDir+File.separator+"conf/examples/1GB/native-standalone/accumulo-env.sh");
        File destFile = new File(accumuloHomeDir+File.separator+"conf/accumulo-env.sh");
        Files.copy(inputFile.toPath(), destFile.toPath(), StandardCopyOption.REPLACE_EXISTING);
    }

    /**
     * Copy native maps from mesosDirectory to accumuloHome/lib/native if the maps exist
     *
     * @param mesosDirectory
     * @param accumuloHome
     * @throws IOException
     */
    public static void copyAccumuloNativeMaps( String mesosDirectory, String accumuloHome) throws IOException {
        File inputFile = new File(mesosDirectory+File.separator+ Constants.ACCUMULO_NATIVE_LIB);
        File destFile = new File(accumuloHome+File.separator+"lib/native");
        if( inputFile.exists() ) {
            Files.copy(inputFile.toPath(), destFile.toPath(), StandardCopyOption.REPLACE_EXISTING);
        }
    }


    private static void logErrorAndDie(String message, Exception e){
        LOGGER.error(message, e);
        throw new RuntimeException(message, e);
    }


} public class Framework  {

  private String bindAddress = Defaults.BIND_ADDRESS;
  private Integer httpPort = Defaults.HTTP_PORT;
  private String mesosMaster = Defaults.MESOS_MASTER;
  private String name;
  private String id;
  private String tarballUri;
  private String dockerImage;
  private String zkServers = Defaults.ZK_SERVERS;
  private Accumulo cluster = null;


  /**
   * IP address to bind framework webserver to.
   **/
  @JsonProperty("bindAddress")
  public String getBindAddress() {
    return bindAddress;
  }
  public void setBindAddress(String bindAddress) {
    this.bindAddress = bindAddress;
  }

  /**
   * Port to bind framework webserver to.
   **/
  @JsonProperty("httpPort")
  public Integer getHttpPort() {
    return httpPort;
  }
  public void setHttpPort(Integer httpPort) {
    this.httpPort = httpPort;
  }

  /**
   * IP and port of Mesos Master node to register with.
   **/
  @JsonProperty("mesosMaster")
  public String getMesosMaster() {
    return mesosMaster;
  }
  public void setMesosMaster(String mesosMaster) {
    this.mesosMaster = mesosMaster;
  }


  /**
   * Descriptive name for this framework.
   **/
  @JsonProperty("name")
  public String getName() {
    return name;
  }
  public void setName(String name) {
    this.name = name;
  }
  public boolean hasName(){ return !Strings.isNullOrEmpty(name); }

  /**
   * Unique ID that references this framework
   **/
  @JsonProperty("id")
  public String getId() {
    return id;
  }
  public void setId(String id) {
    this.id = id;
  }
  public boolean hasId(){ return !Strings.isNullOrEmpty(id); }

  /**
   * URI of tarball containing framework distribution
   **/
  @JsonProperty("tarballUri")
  public String getTarballUri() {
    return tarballUri;
  }
  public void setTarballUri(String tarballUri) {
    this.tarballUri = tarballUri;
  }
  public boolean hasTarballUri(){ return !Strings.isNullOrEmpty(tarballUri);  }

  /**
  * Name of docker container to use for deployment
  *
  * @return
  */
  @JsonProperty("dockerInfo")
  public String getDockerImage() { return dockerImage; }
  public void setDockerImage(String dockerImage){ this.dockerImage = dockerImage; }
  public boolean hasDockerImage() { return !Strings.isNullOrEmpty(dockerImage);  }

  /**
   * List of Zookeeper servers to store Framework state on. Does not have\nto be the same ZK servers used for the Accumulo cluster
   **/
  @JsonProperty("zkServers")
  public String getZkServers() {
    return zkServers;
  }
  public void setZkServers(String zkServers) {
    this.zkServers = zkServers;
  }


  /**
   * Accumulo cluster definition
   **/
  @JsonProperty("cluster")
  public Accumulo getCluster() {
    return cluster;
  }
  public void setCluster(Accumulo cluster) {
    this.cluster = cluster;
  }
  public boolean hasCluster(){
    return isNotNullEmpty(cluster);
  }

  private static boolean isNotNullEmpty(Object thing){
    boolean hasThing = false;
    if( thing != null ){
      hasThing = true;
      if( thing instanceof String){
        if( ((String) thing).trim().isEmpty() ){
          hasThing = false;
        }
      }
    }
    return hasThing;
  }

    public void merge(Framework other){
        if(!this.hasName() && other.hasName()) this.setName(other.getName());
        if(!this.hasId() && other.hasId()) this.setId(other.getId());
        if(!this.hasDockerImage() && other.hasDockerImage()) this.setDockerImage(other.getDockerImage());
        if(!this.hasTarballUri() && other.hasTarballUri()) this.setTarballUri(other.getTarballUri());
        if(!this.hasCluster() && other.hasCluster()) this.setCluster(other.getCluster());
    }

  @Override
  public String toString()  {

    MoreObjects.ToStringHelper helper = MoreObjects.toStringHelper(this);
            helper
            .add("bindAddress", bindAddress)
            .add("httpPort", httpPort)
            .add("dockerImage", this.dockerImage)
            .add("id", this.id)
            .add("mesosMaster", this.mesosMaster)
            .add("name", this.name)
            .add("tarballUri", this.tarballUri)
            .add("zkServers", this.zkServers);
        if( this.hasCluster() ){
            helper.add("cluster", this.cluster.toString());
        } else {
            helper.add("cluster", "");
        }

    return helper.toString();

  }

  @Override
  public int hashCode(){
    return Objects.hash(bindAddress,
            httpPort,
            cluster,
            dockerImage,
            id,
            mesosMaster,
            name,
            tarballUri,
            zkServers);
  }
} public enum Cluster {
    INSTANCE;

    private static final Logger LOGGER = LoggerFactory.getLogger(Cluster.class);

    private volatile Framework config;
    private State state;

    private volatile List<Task> tasks = new ArrayList<>();
    private volatile Set<Protos.TaskStatus> runningTasks = Sets.newConcurrentHashSet();

    private Matcher matcher;
    private Launcher launcher;

    // enum constructors are private by definition. This is where some static initializations could go.
    Cluster(){}

    /**
        Initializes the cluster state from the config.

        Must be called before running Cluster with Mesos.
     */
    public void initialize(Framework config, State state){
        this.config = config;
        this.state = state;

        // TODO consider singleton launchers and matchers defined per server group
        this.launcher = new AccumuloStartExecutorLauncher(config);
        this.matcher = new MinCpuMinRamFIFOMatcher(config.getCluster());

        initializeTasks();
    }

    private void initializeTasks(){
        Accumulo cluster = config.getCluster();
        List<ServerGroup> serverGroups = cluster.getServerGroups();
        for( ServerGroup group : serverGroups){
            for( int ii = 0; ii < group.getCount(); ii++){
                Task task = new Task();
                task.setServerProfile(group.getProfile());
                tasks.add(task);
            }
        }
    }

    public void setFrameworkId(String fid){
        config.setId(fid);
    }

    public void handleOffers(SchedulerDriver driver, List<Protos.Offer> offers){
        int originalOfferSize = offers.size();

        LOGGER.info("Received {} offers", originalOfferSize);
        List<Task> tasksToLaunch = this.getTasksToLaunch();
        if( !tasksToLaunch.isEmpty() ) {
            //TODO right now matchOffers removes offers it uses from the offers collection
            // this assumed side-effect is bad if we actually want people to write their own
            // matchers.
            List<Match> matchedServers = matcher.matchOffers(tasksToLaunch, offers);

            LOGGER.info("Found {} matches for servers from {} offers",
                    matchedServers.size(), originalOfferSize);

            // Launch all the matched servers.
            for (Match match : matchedServers) {
                LOGGER.info("Launching Server: {} using offer {}", match.getTask().getType().name(),
                        match.getOffer().getId());

                match.getTask().assignTaskId();
                Protos.TaskInfo taskInfo = launcher.launch(driver, match);
                match.getTask().setSlaveId(taskInfo.getSlaveId().getValue());

                LOGGER.info("Created Task {} on slave {}", taskInfo.getTaskId(), taskInfo.getSlaveId().getValue());
            }

        } else {
            LOGGER.info("No tasks currently waiting launch.");
        }
        LOGGER.info("Declining {} offers", offers.size());
        declineUnmatchedOffers(driver, offers);
    }

    /**
     * Iterates through tasks finding ones without taskId or slaveId
     *
     * @return List of unlaunched tasks
     */
    private List<Task> getTasksToLaunch(){
        List<Task> toLaunch = Lists.newArrayList();
        for( Task task: tasks){
            if( !task.hasTaskId() && !task.hasSlaveId() ){
                toLaunch.add(task);
            }
        }
        return toLaunch;
    }

    public void restore(SchedulerDriver driver) {

        // reconcileTasks causes the framework to call updateTaskStatus, which
        // will update the tasks list.
        // TODO handle return of reconcileTasks
        Set<Protos.TaskStatus> emptyStatus = Sets.newHashSet();
        Protos.Status reconcileStatus = driver.reconcileTasks(emptyStatus);

        //TODO save cluster state?
    }

    /**
     * Updates Cluster state based on task status.
     *
     * @param status
     */
    public void updateTaskStatus(Protos.TaskStatus status){

        String slaveId = status.getSlaveId().getValue();
        String taskId = status.getTaskId().getValue();


        Task task = findTask(taskId, slaveId);

        if( task == null ){
            // TODO kill task? illegal state exception?
        }

        LOGGER.info("Task Status Update: Status: {} Slave: {} Task: {}", status.getState(), slaveId, taskId);

        switch (status.getState()){
            case TASK_RUNNING:
                runningTasks.add(status);
                break;
            case TASK_FINISHED:
            case TASK_FAILED:
            case TASK_KILLED:
            case TASK_LOST:
                runningTasks.remove(status);
                // Reset the task, so it will be re-launched
                synchronized(task) {
                    task.setSlaveId(null);
                    task.setTaskId(null);
                }
                
                break;
            // TODO task error
            case TASK_STARTING:
            case TASK_STAGING:
                break;
            default:
                LOGGER.info("Unknown Task Status received: {}", status.getState().toString());

        }
    }

    // Finds a Task based on taskId and slaveId
    private Task findTask(@NotNull String taskId, @NotNull String slaveId){
        LOGGER.info("Finding task {} slaveId {}", taskId, slaveId);
        for( Task task : tasks ){
            if( task.hasTaskId() && task.getTaskId().equals(taskId)
                    && task.hasSlaveId() && task.getSlaveId().equals(slaveId)){
                return task;
            }
        }
        return null;
    }

    // Decline any offers not claimed by a match
    private void declineUnmatchedOffers(SchedulerDriver driver, List<Protos.Offer> offers){
        for (Protos.Offer offer : offers){
            LOGGER.debug("Declining Offer: ", offer.getId().getValue());
            driver.declineOffer(offer.getId());
        }
    }

} public class Scheduler implements org.apache.mesos.Scheduler {
    private static final Logger LOGGER = LoggerFactory.getLogger(Scheduler.class);

    private final Cluster cluster;

    public Scheduler(final Cluster cluster){
        this.cluster = cluster;
    }

    /**
     * Invoked when the scheduler successfully registers with a Mesos
     * master. A unique ID (generated by the master) used for
     * distinguishing this framework from others and MasterInfo
     * with the IP and port of the current master are provided as arguments.
     *
     * @param schedulerDriver
     * @param frameworkID
     * @param masterInfo      Info about the current master, including IP and port.
     * @see org.apache.mesos.SchedulerDriver
     * @see Protos.FrameworkID
     * @see Protos.MasterInfo
     */
    @Override
    public void registered(SchedulerDriver schedulerDriver, Protos.FrameworkID frameworkID, Protos.MasterInfo masterInfo) {
        LOGGER.info("Framework Registered: {}", frameworkID.getValue());
        cluster.setFrameworkId(frameworkID.getValue()); // ID is only available after registration.
        cluster.restore(schedulerDriver);
    }

    /**
     * Invoked when the scheduler re-registers with a newly elected Mesos master.
     * This is only called when the scheduler has previously been registered.
     * MasterInfo containing the updated information about the elected master
     * is provided as an argument.
     *
     * @param schedulerDriver
     * @param masterInfo      The updated information about the elected master.
     * @see org.apache.mesos.SchedulerDriver
     * @see Protos.MasterInfo
     */
    @Override
    public void reregistered(SchedulerDriver schedulerDriver, Protos.MasterInfo masterInfo) {
        LOGGER.info("Framework Re-Registered");
        cluster.restore(schedulerDriver);
    }

    /**
     * Invoked when resources have been offered to this framework. A
     * single offer will only contain resources from a single slave.
     * Resources associated with an offer will not be re-offered to
     * _this_ framework until either (a) this framework has rejected
     * those resources (see {@link org.apache.mesos.SchedulerDriver#launchTasks}) or (b)
     * those resources have been rescinded (see {@link org.apache.mesos.Scheduler#offerRescinded}).
     * Note that resources may be concurrently offered to more than one
     * framework at a time (depending on the allocator being used). In
     * that case, the first framework to launch tasks using those
     * resources will be able to use them while the other frameworks
     * will have those resources rescinded (or if a framework has
     * already launched tasks with those resources then those tasks will
     * fail with a TASK_LOST status and a message saying as much).
     *
     * @param schedulerDriver
     * @param list
     * @see org.apache.mesos.SchedulerDriver
     * @see Protos.Offer
     */
    @Override
    public void resourceOffers(SchedulerDriver schedulerDriver, List<Protos.Offer> offers) {
        LOGGER.info("{} Resource Offers Recieved", offers.size());
        cluster.handleOffers(schedulerDriver, offers);
    }

    /**
     * Invoked when an offer is no longer valid (e.g., the slave was
     * lost or another framework used resources in the offer). If for
     * whatever reason an offer is never rescinded (e.g., dropped
     * message, failing over framework, etc.), a framwork that attempts
     * to launch tasks using an invalid offer will receive TASK_LOST
     * status updats for those tasks (see {@link #resourceOffers}).
     *
     * @param schedulerDriver
     * @param offerID
     * @see org.apache.mesos.SchedulerDriver
     * @see Protos.OfferID
     */
    @Override
    public void offerRescinded(SchedulerDriver schedulerDriver, Protos.OfferID offerID) {
        LOGGER.info("Offer Rescinded");
    }

    /**
     * Invoked when the status of a task has changed (e.g., a slave is
     * lost and so the task is lost, a task finishes and an executor
     * sends a status update saying so, etc). If implicit
     * acknowledgements are being used, then returning from this
     * callback _acknowledges_ receipt of this status update! If for
     * whatever reason the scheduler aborts during this callback (or
     * the process exits) another status update will be delivered (note,
     * however, that this is currently not true if the slave sending the
     * status update is lost/fails during that time). If explicit
     * acknowledgements are in use, the scheduler must acknowledge this
     * status on the driver.
     *
     * @param schedulerDriver
     * @param taskStatus
     * @see org.apache.mesos.SchedulerDriver
     * @see Protos.TaskStatus
     */
    @Override
    public void statusUpdate(SchedulerDriver schedulerDriver, Protos.TaskStatus taskStatus) {
        LOGGER.info("Status Update Received");
        cluster.updateTaskStatus(taskStatus);
    }

    /**
     * Invoked when an executor sends a message. These messages are best
     * effort; do not expect a framework message to be retransmitted in
     * any reliable fashion.
     *
     * @param schedulerDriver
     * @param executorID
     * @param slaveID
     * @param bytes           @see org.apache.mesos.SchedulerDriver
     * @see Protos.ExecutorID
     * @see Protos.SlaveID
     */
    @Override
    public void frameworkMessage(SchedulerDriver schedulerDriver, Protos.ExecutorID executorID, Protos.SlaveID slaveID, byte[] bytes) {
        LOGGER.info("Framework Message received.");
    }

    /**
     * Invoked when the scheduler becomes "disconnected" from the master
     * (e.g., the master fails and another is taking over).
     *
     * @param schedulerDriver@see org.apache.mesos.SchedulerDriver
     */
    @Override
    public void disconnected(SchedulerDriver schedulerDriver) {
        LOGGER.info("Disconnected from Mesos master");
        // new scheduler should reconnect us, but perhaps set something to do after a timeout here.
    }

    /**
     * Invoked when a slave has been determined unreachable (e.g.,
     * machine failure, network partition). Most frameworks will need to
     * reschedule any tasks launched on this slave on a new slave.
     *
     * @param schedulerDriver
     * @param slaveID
     * @see org.apache.mesos.SchedulerDriver
     * @see Protos.SlaveID
     */
    @Override
    public void slaveLost(SchedulerDriver schedulerDriver, Protos.SlaveID slaveID) {
        LOGGER.info("Slave Lost: " + slaveID);
        // TODO notify cluster that a slave was lost. Have executor kill if still running
    }

    /**
     * Invoked when an executor has exited/terminated. Note that any
     * tasks running will have TASK_LOST status updates automagically
     * generated.
     *
     * @param schedulerDriver
     * @param executorID
     * @param slaveID
     * @param i               @see org.apache.mesos.SchedulerDriver
     * @see Protos.ExecutorID
     * @see Protos.SlaveID
     * @see Protos.Status
     */
    @Override
    public void executorLost(SchedulerDriver schedulerDriver, Protos.ExecutorID executorID, Protos.SlaveID slaveID, int i) {
        LOGGER.info("Executor lost: " + executorID + "on slave: " + slaveID + " number " + i);
        // Will update state when TASK_LOST received
    }

    /**
     * Invoked when there is an unrecoverable error in the scheduler or
     * driver. The driver will be aborted BEFORE invoking this callback.
     *
     * @param schedulerDriver
     * @param s
     * @see org.apache.mesos.SchedulerDriver
     */
    @Override
    public void error(SchedulerDriver schedulerDriver, String s) {
        LOGGER.info("Unrecoverable Error in Scheduler or SchedulerDriver: " + s);
    }
} public class FrameworkStateHelper {
    private static final Logger LOGGER = LoggerFactory.getLogger(FrameworkStateHelper.class);

    private static final ObjectMapper mapper = new ObjectMapper();
    private static final String REGISTERED_FRAMEWORKS_KEY = "registered_frameworks";
    private static final String FRAMEWORK_CONFIG_PREFIX = "framework__";

    private final State state;

    public FrameworkStateHelper(State state){
        this.state = state;
    }

    /**
     * Checks if any frameworks have been written to this state interface.
     *
     * @return
     * @throws InterruptedException
     */
    public boolean hasRegisteredFrameworks() throws InterruptedException {
        return keyExists(REGISTERED_FRAMEWORKS_KEY);
    }

    private IdRegistry getFrameworkRegistry() throws ExecutionException, InterruptedException, IOException {
        IdRegistry registry;
        if( hasRegisteredFrameworks() ) {
            byte[] registryBytes = getBytesForKey(REGISTERED_FRAMEWORKS_KEY);
            registry = mapper.readValue(registryBytes, IdRegistry.class);
        } else {
            registry = new IdRegistry();
        }
        return registry;
    }

    public Map<String, String> getFrameworkIdMap() throws ExecutionException, InterruptedException, IOException {
        IdRegistry registry = getFrameworkRegistry();
        Map<String,String> retMap = Maps.newHashMap();
        for(IdRegistry.RegistryPair pair : registry.getRegistry()){
            retMap.put(pair.getId(), pair.getName());
        }
        return retMap;
    }

    /**
     * Retrieves a map of framework name:id pairs
     *
     * @return
     * @throws ExecutionException
     * @throws InterruptedException
     * @throws IOException
     */
    public Map<String, String> getFrameworkNameMap() throws ExecutionException, InterruptedException, IOException {
        IdRegistry registry = getFrameworkRegistry();
        Map<String,String> retMap = Maps.newHashMap();
        for(IdRegistry.RegistryPair pair : registry.getRegistry()){
            retMap.put(pair.getName(), pair.getId());
        }
        return retMap;
    }

    public void saveFrameworkConfig(Framework config) throws IOException, ExecutionException, InterruptedException {

        byte[] configBytes = mapper.writeValueAsBytes(config);
        if( config.hasId() ) {
            String key = buildFrameworkConfigKey(config.getId());
            LOGGER.info("Saving framework config: {} bytes at {}", configBytes.length, key);
            saveBytesForKey(key, configBytes);
        } else {
            throw new ExecutionException(new Exception("Cannot save configuration with no defined id"));
        }
        updateRegistry(config);
    }

    public Framework getFrameworkConfig(String id) throws ExecutionException, InterruptedException, IOException {
        String key = buildFrameworkConfigKey(id);
        byte[] configBytes = getBytesForKey(key);
        LOGGER.info("Read config: {} bytes at {}", configBytes.length, key);
        Framework config = mapper.readValue(configBytes, Framework.class);

        return config;
    }

    private void updateRegistry(Framework config) throws InterruptedException, ExecutionException, IOException {
        synchronized (REGISTERED_FRAMEWORKS_KEY){
            boolean foundFramework = false;
            List<IdRegistry.RegistryPair> newRegistry = Lists.newArrayList();
            IdRegistry registry;
            if( hasRegisteredFrameworks() ) {
                registry = getFrameworkRegistry();
                for (IdRegistry.RegistryPair pair : registry.getRegistry()) {
                    if (pair.getId().equals(config.getId())) {
                        foundFramework = true;
                    }
                    newRegistry.add(pair);
                }
            } else {
                registry = new IdRegistry();
            }
            if( !foundFramework ){
                newRegistry.add(new IdRegistry.RegistryPair(config.getName(), config.getId()));
            }
            registry.setRegistry(newRegistry);
            byte[] regBytes = mapper.writeValueAsBytes(registry);
            saveBytesForKey(REGISTERED_FRAMEWORKS_KEY, regBytes);
        }
    }

    private boolean keyExists(String inKey) throws InterruptedException {
        boolean exists = false;
        try {
            Future<Iterator<String>> keys = state.names();
            for (Iterator<String> iter = keys.get(); iter.hasNext(); ) {
                String key = iter.next();
                if (key.equals(inKey)) {
                    exists = true;
                    break;
                }
            }
        } catch (ExecutionException e) {
            LOGGER.error("Key not found via Execution Exception : {}", inKey);
            // key doesn't exist
        }
        return exists;
    }

    private byte[] getBytesForKey(String key) throws InterruptedException, ExecutionException {
        Variable stateVariable = getVariableForKey(key);
        return stateVariable.value();
    }

    private Variable getVariableForKey(String key) throws ExecutionException, InterruptedException {
        Future<Variable> stateFuture = state.fetch(key);
        Variable stateVariable = stateFuture.get();

        return stateVariable;
    }

    private void saveBytesForKey(String key, byte[] bytes) throws InterruptedException, ExecutionException {
        Variable variable;
        variable = getVariableForKey(key);
        Variable toSave = variable.mutate(bytes);
        this.state.store(toSave);
    }

    private static String buildFrameworkConfigKey(String id){
        return FRAMEWORK_CONFIG_PREFIX + id;
    }
}
import aredee.mesos.frameworks.accumulo.configuration.CommandLineHandler;
import aredee.mesos.frameworks.accumulo.configuration.Constants;
import aredee.mesos.frameworks.accumulo.configuration.Defaults;
import aredee.mesos.frameworks.accumulo.configuration.Environment;
import aredee.mesos.frameworks.accumulo.framework.api.WebServer;
import aredee.mesos.frameworks.accumulo.framework.guice.ApiServletModule;
import aredee.mesos.frameworks.accumulo.framework.guice.ConfigurationModule;
import aredee.mesos.frameworks.accumulo.initialize.AccumuloInitializer;
import aredee.mesos.frameworks.accumulo.model.Framework;
import aredee.mesos.frameworks.accumulo.scheduler.Cluster;
import aredee.mesos.frameworks.accumulo.scheduler.Scheduler;
import aredee.mesos.frameworks.accumulo.state.FrameworkStateHelper;
import com.google.common.base.Preconditions;
import com.google.common.base.Strings;
import com.google.inject.Guice;
import com.google.inject.Injector;
import com.google.inject.Stage;
import org.apache.mesos.MesosSchedulerDriver;
import org.apache.mesos.Protos.FrameworkInfo;
import org.apache.mesos.SchedulerDriver;
import org.apache.mesos.state.State;
import org.apache.mesos.state.ZooKeeperState;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.ExecutionException;
/**
 *
 * 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 aredee.mesos.frameworks.accumulo.framework;



public final class Main {

    private static final Logger LOGGER = LoggerFactory.getLogger(Main.class);

    private WebServer webServer;

    public static void main(String[] args) throws IOException, ExecutionException, InterruptedException {


        CommandLineHandler cmdHandler = new CommandLineHandler(args);
        if( cmdHandler.checkHelpOrVersion() ){
            System.exit(0);  // checkHelpOrVersion prints appropriate info to System.out
        }

        if( !environmentVariablesAreSet() ){
            System.exit(-1);
        }

        Framework frameworkConfig = cmdHandler.getFrameworkDefinition();

        LOGGER.info("Connecting to Mesos state");
        State state = getMesosState(frameworkConfig);
        FrameworkStateHelper stateHelper = new FrameworkStateHelper(state);

        // at this point perhaps only the framework name was given, so
        // read framework config
        String frameworkName = frameworkConfig.getName();
        Map<String,String> frameworks = stateHelper.getFrameworkNameMap();
        LOGGER.info("frameworks ? {}", frameworks);
        if( frameworks.containsKey(frameworkName) ){
            LOGGER.info("Found previous framework with name {}", frameworkName);

            Framework savedConfig = stateHelper.getFrameworkConfig(frameworks.get(frameworkName));
            LOGGER.info("Found saved framework config: {}", savedConfig );
            frameworkConfig.merge(savedConfig);
            LOGGER.info("Merged command line configuration with saved configuration. Command line trumps.");
        }

        // check if this is an initialization request
        if( cmdHandler.isInitializeRequest() ){

            LOGGER.info("Initializing Accumulo Instance: {}", frameworkConfig.getCluster().getInstance());
            AccumuloInitializer accumuloInitializer =
                    new AccumuloInitializer(frameworkConfig.getCluster());

            int status = accumuloInitializer.initialize();
            if( status != 0 ){
                LOGGER.error("Accumulo failed to initialize instance: {}", frameworkConfig.getCluster().getInstance());
                LOGGER.error("Either remove instance from HDFS or re-launch framework without init option");
                System.exit(status);
            } else {
                LOGGER.info("Accumulo instance initialized: {}", frameworkConfig.getCluster().getInstance());
            }

            // after initialization, save the config
            frameworkConfig.setId(UUID.randomUUID().toString());

            try {
                stateHelper.saveFrameworkConfig(frameworkConfig);
                LOGGER.info("Saved Framework Configuration: {}", frameworkConfig.getId());

            } catch (Exception e) {
                LOGGER.error("Unable to save framework configuration to Mesos State");
                System.err.println("Unable to save framework configuration to Mesos State");
                e.printStackTrace();
                System.exit(1);
            }
        } else {

            int exitStatus = -1;
            try {
                LOGGER.info("Starting mesos-accumumlo framework version " + Constants.FRAMEWORK_VERSION);
                exitStatus = new Main().run(frameworkConfig, state);
            } catch (Exception e) {
                LOGGER.error("Unhandled exception encountered, exiting: ", e);
            }

            LOGGER.info("Exiting Accumulo Framework with status: " + exitStatus);
            System.exit(exitStatus);
        }

        System.exit(0); // necessary for init to finish because state launches a process
    }

    private int run(Framework config, State mesosState) throws Exception {

        // Sanity check before firing up processes.
        LOGGER.info("Preconditions check: {}", config);

        Preconditions.checkState(!Strings.isNullOrEmpty(config.getId()));
        Preconditions.checkState( !Strings.isNullOrEmpty(config.getName()));
        Preconditions.checkState( config.hasCluster(), "No cluster definition found, exiting" );

        // Initialize the Cluster singleton. Do this before creating the webservice in case the
        // webservice touches it.
        Cluster cluster = Cluster.INSTANCE;
        cluster.initialize(config, mesosState);

        Injector injector = Guice.createInjector(
                Stage.PRODUCTION,
                new ConfigurationModule(config),
                new ApiServletModule()
        );

        startWebserverOrDie(injector);
        
        // Start the schedulerDriver
        LOGGER.info("Initializing mesos-accumulo Scheduler");
  
        FrameworkInfo frameworkInfo = FrameworkInfo.newBuilder()
                .setId(createMesosFrameworkID(config.getId())) // empty string creates new random name
                .setName(config.getName())
                .setCheckpoint(true)
                .setHostname("") // let mesos set to current hostname
                .setUser("")  // empty string is current user of mesos process
                .setWebuiUrl("http://" + config.getBindAddress() + ":" + config.getHttpPort() + "/")
                .build();

        final SchedulerDriver schedulerDriver =
                new MesosSchedulerDriver( new Scheduler(cluster), frameworkInfo, config.getMesosMaster());

        LOGGER.info("Running mesos-accumulo SchedulerDriver");
        final int status;
        
        // run() blocks until driver finishes somehow
        org.apache.mesos.Protos.Status driverStatus = schedulerDriver.run();
        
        switch (driverStatus) {
            case DRIVER_STOPPED:
                status = 0;
                break;
            case DRIVER_ABORTED:
                status = 1;
                break;
            case DRIVER_NOT_STARTED:
                status = 2;
                break;
            default:
                status = 3;
                break;
        }
        LOGGER.info("mesos-accumulo stopped with status " +
                driverStatus.name() + " " + driverStatus.getNumber());

        webServer.stop();

        // Ensure that the driver process terminates.
        schedulerDriver.stop(true);

        return status;
    }

    private static State getMesosState(Framework config){
        return new ZooKeeperState(config.getZkServers(),
                Defaults.ZOOKEEPER_TIMEOUT,