package org.apache.helix.tools;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.InputStream;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

import org.apache.helix.HelixAdmin;
import org.apache.helix.HelixException;
import org.apache.helix.manager.zk.ZKHelixAdmin;
import org.apache.helix.manager.zk.ZKHelixManager;
import org.apache.helix.model.HelixConfigScope;
import org.apache.helix.model.HelixConfigScope.ConfigScopeProperty;
import org.apache.helix.model.IdealState;
import org.apache.helix.model.IdealState.RebalanceMode;
import org.apache.helix.model.InstanceConfig;
import org.apache.helix.model.StateModelDefinition;
import org.apache.helix.model.builder.HelixConfigScopeBuilder;
import org.apache.helix.tools.YAMLClusterSetup.YAMLClusterConfig.ParticipantConfig;
import org.apache.helix.tools.YAMLClusterSetup.YAMLClusterConfig.ResourceConfig;
import org.apache.helix.tools.YAMLClusterSetup.YAMLClusterConfig.ResourceConfig.ConstraintsConfig;
import org.apache.helix.tools.YAMLClusterSetup.YAMLClusterConfig.ResourceConfig.StateModelConfig;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.yaml.snakeyaml.Yaml;

/*
 * 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.
 */

/**
 * Supports HelixAdmin operations specified by a YAML configuration file defining a cluster,
 * resources, participants, etc.
 * See the user-rebalanced-lock-manager recipe for an annotated example file.
 */
public class YAMLClusterSetup {
  private static final Logger LOG = LoggerFactory.getLogger(YAMLClusterSetup.class);

  private final String _zkAddress;

  /**
   * Start the YAML parser for a given zookeeper instance
   * @param zkAddress
   */
  public YAMLClusterSetup(String zkAddress) {
    _zkAddress = zkAddress;
  }

  /**
   * Set up the cluster by parsing a YAML file.
   * @param input InputStream representing the file
   * @return ClusterConfig Java wrapper of the configuration file
   */
  public YAMLClusterConfig setupCluster(InputStream input) {
    // parse the YAML
    Yaml yaml = new Yaml();
    YAMLClusterConfig cfg = yaml.loadAs(input, YAMLClusterConfig.class);

    // create the cluster
    HelixAdmin helixAdmin = new ZKHelixAdmin(_zkAddress);
    if (cfg.clusterName == null) {
      throw new HelixException("Cluster name is required!");
    }
    helixAdmin.addCluster(cfg.clusterName);

    // add each participant
    if (cfg.participants != null) {
      for (ParticipantConfig participant : cfg.participants) {
        helixAdmin.addInstance(cfg.clusterName, getInstanceCfg(participant));
      }
    }

    // add each resource
    if (cfg.resources != null) {
      for (ResourceConfig resource : cfg.resources) {
        if (resource.name == null) {
          throw new HelixException("Resources must be named!");
        }
        if (resource.stateModel == null || resource.stateModel.name == null) {
          throw new HelixException("Resource must specify a named state model!");
        }
        // if states is null, assume using a built-in or already-added state model
        if (resource.stateModel.states != null) {
          StateModelDefinition stateModelDef =
              getStateModelDef(resource.stateModel, resource.constraints);
          helixAdmin.addStateModelDef(cfg.clusterName, resource.stateModel.name, stateModelDef);
        } else {
          StateModelDefinition stateModelDef = null;
          if (resource.stateModel.name.equals("MasterSlave")) {
            stateModelDef =
                new StateModelDefinition(StateModelConfigGenerator.generateConfigForMasterSlave());
          } else if (resource.stateModel.name.equals("OnlineOffline")) {
            stateModelDef =
                new StateModelDefinition(StateModelConfigGenerator.generateConfigForOnlineOffline());
          } else if (resource.stateModel.name.equals("LeaderStandby")) {
            stateModelDef =
                new StateModelDefinition(StateModelConfigGenerator.generateConfigForLeaderStandby());
          }
          if (stateModelDef != null) {
            try {
              helixAdmin.addStateModelDef(cfg.clusterName, resource.stateModel.name, stateModelDef);
            } catch (HelixException e) {
              LOG.warn("State model definition " + resource.stateModel.name
                  + " could not be added.");
            }
          }
        }
        int partitions = 1;
        int replicas = 1;
        if (resource.partitions != null) {
          if (resource.partitions.containsKey("count")) {
            partitions = resource.partitions.get("count");
          }
          if (resource.partitions.containsKey("replicas")) {
            replicas = resource.partitions.get("replicas");
          }
        }

        if (resource.rebalancer == null || !resource.rebalancer.containsKey("mode")) {
          throw new HelixException("Rebalance mode is required!");
        }
        helixAdmin.addResource(cfg.clusterName, resource.name, partitions,
            resource.stateModel.name, resource.rebalancer.get("mode"));
        // user-defined rebalancer
        if (resource.rebalancer.containsKey("class")
            && resource.rebalancer.get("mode").equals(RebalanceMode.USER_DEFINED.toString())) {
          IdealState idealState = helixAdmin.getResourceIdealState(cfg.clusterName, resource.name);
          idealState.setRebalancerClassName(resource.rebalancer.get("class"));
          helixAdmin.setResourceIdealState(cfg.clusterName, resource.name, idealState);
        }
        helixAdmin.rebalance(cfg.clusterName, resource.name, replicas);
      }
    }

    // enable auto join if this option is set
    if (cfg.autoJoinAllowed != null && cfg.autoJoinAllowed) {
      HelixConfigScope scope =
          new HelixConfigScopeBuilder(ConfigScopeProperty.CLUSTER).forCluster(cfg.clusterName)
              .build();
      Map<String, String> properties = new HashMap<String, String>();
      properties.put(ZKHelixManager.ALLOW_PARTICIPANT_AUTO_JOIN, cfg.autoJoinAllowed.toString());
      helixAdmin.setConfig(scope, properties);
    }
    return cfg;
  }

  private static InstanceConfig getInstanceCfg(ParticipantConfig participant) {
    if (participant == null || participant.name == null || participant.host == null
        || participant.port == null) {
      throw new HelixException("Participant must have a specified name, host, and port!");
    }
    InstanceConfig instanceCfg = new InstanceConfig(participant.name);
    instanceCfg.setHostName(participant.host);
    instanceCfg.setPort(participant.port.toString());
    return instanceCfg;
  }

  private static StateModelDefinition getStateModelDef(StateModelConfig stateModel,
      ConstraintsConfig constraints) {
    // Use a builder to define the state model
    StateModelDefinition.Builder builder = new StateModelDefinition.Builder(stateModel.name);
    if (stateModel.states == null || stateModel.states.size() == 0) {
      throw new HelixException("List of states are required in a state model!");
    }
    Set<String> stateSet = new HashSet<String>(stateModel.states);
    if (stateModel.initialState == null) {
      throw new HelixException("Initial state is required in a state model!");
    } else if (!stateSet.contains(stateModel.initialState)) {
      throw new HelixException("Initial state is not a valid state");
    }
    builder.initialState(stateModel.initialState);

    // Build a helper for state priorities
    Map<String, Integer> statePriorities = new HashMap<String, Integer>();
    if (constraints != null && constraints.state != null && constraints.state.priorityList != null) {
      int statePriority = 0;
      for (String state : constraints.state.priorityList) {
        if (!stateSet.contains(state)) {
          throw new HelixException("State " + state
              + " in the state priority list is not in the state list!");
        }
        statePriorities.put(state, statePriority);
        statePriority++;
      }
    }

    // Add states, set state priorities
    for (String state : stateModel.states) {
      if (statePriorities.containsKey(state)) {
        builder.addState(state, statePriorities.get(state));
      } else {
        builder.addState(state);
      }
    }

    // Set state counts
    for (Map<String, String> counts : constraints.state.counts) {
      String state = counts.get("name");
      if (!stateSet.contains(state)) {
        throw new HelixException("State " + state + " has a count, but not in the state list!");
      }
      builder.dynamicUpperBound(state, counts.get("count"));
    }

    // Build a helper for transition priorities
    Map<String, Integer> transitionPriorities = new HashMap<String, Integer>();
    if (constraints != null && constraints.transition != null
        && constraints.transition.priorityList != null) {
      int transitionPriority = 0;
      for (String transition : constraints.transition.priorityList) {
        transitionPriorities.put(transition, transitionPriority);
        transitionPriority++;
      }
    }

    // Add the transitions
    if (stateModel.transitions == null || stateModel.transitions.size() == 0) {
      throw new HelixException("Transitions are required!");
    }
    for (Map<String, String> transitions : stateModel.transitions) {
      String name = transitions.get("name");
      String from = transitions.get("from");
      String to = transitions.get("to");
      if (name == null || from == null || to == null) {
        throw new HelixException("All transitions must have a name, a from state, and a to state");
      }
      if (transitionPriorities.containsKey(name)) {
        builder.addTransition(from, to, transitionPriorities.get(name));
      } else {
        builder.addTransition(from, to);
      }
    }

    return builder.build();
  }

  /**
   * Java wrapper for the YAML input file
   */
  public static class YAMLClusterConfig {
    public String clusterName;
    public List<ResourceConfig> resources;
    public List<ParticipantConfig> participants;
    public Boolean autoJoinAllowed;

    public static class ResourceConfig {
      public String name;
      public Map<String, String> rebalancer;
      public Map<String, Integer> partitions;
      public StateModelConfig stateModel;
      public ConstraintsConfig constraints;

      public static class StateModelConfig {
        public String name;
        public List<String> states;
        public List<Map<String, String>> transitions;
        public String initialState;
      }

      public static class ConstraintsConfig {
        public StateConstraintsConfig state;
        public TransitionConstraintsConfig transition;

        public static class StateConstraintsConfig {
          public List<Map<String, String>> counts;
          public List<String> priorityList;
        }

        public static class TransitionConstraintsConfig {
          public List<String> priorityList;
        }
      }
    }

    public static class ParticipantConfig {
      public String name;
      public String host;
      public Integer port;
    }
  }

  /**
   * Start a cluster defined by a YAML file
   * @param args zkAddr, yamlFile
   */
  public static void main(String[] args) {
    if (args.length < 2) {
      LOG.error("USAGE: YAMLClusterSetup zkAddr yamlFile");
      return;
    }
    String zkAddress = args[0];
    String yamlFile = args[1];

    InputStream input;
    try {
      input = new FileInputStream(new File(yamlFile));
    } catch (FileNotFoundException e) {
      LOG.error("Could not open " + yamlFile);
      return;
    }
    new YAMLClusterSetup(zkAddress).setupCluster(input);
  }
}
