/*
 * 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.
 */
package org.apache.ambari.server.stack.upgrade.orchestrate;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;

import org.apache.ambari.server.stack.HostsType;
import org.apache.ambari.server.stack.upgrade.ExecuteHostType;
import org.apache.ambari.server.stack.upgrade.ExecuteTask;
import org.apache.ambari.server.stack.upgrade.Grouping;
import org.apache.ambari.server.stack.upgrade.ParallelClientGrouping;
import org.apache.ambari.server.stack.upgrade.ServiceCheckGrouping.ServiceCheckStageWrapper;
import org.apache.ambari.server.stack.upgrade.Task;
import org.apache.ambari.server.stack.upgrade.Task.Type;
import org.apache.ambari.server.stack.upgrade.UpgradePack.ProcessingComponent;
import org.apache.ambari.server.utils.StageUtils;
import org.apache.commons.collections4.CollectionUtils;

/**
 * Responsible for building the stages for {@link ParallelClientGrouping}
 */
public class ParallelClientGroupingBuilder extends StageWrapperBuilder {

  private Map<String, HostHolder> serviceToHostMap = new HashMap<>();

  /**
   * @param grouping
   */
  public ParallelClientGroupingBuilder(Grouping grouping) {
    super(grouping);
  }

  @Override
  public void add(UpgradeContext upgradeContext, HostsType hostsType, String service,
      boolean clientOnly, ProcessingComponent pc, Map<String, String> params) {

    if (null == hostsType || CollectionUtils.isEmpty(hostsType.getHosts())) {
      return;
    }

    Task task = resolveTask(upgradeContext, pc);
    // !!! better not
    if (null == task) {
      return;
    }

    Iterator<String> hostIterator = hostsType.getHosts().iterator();
    HostHolder holder = new HostHolder();
    holder.m_firstHost = hostIterator.next();

    while (hostIterator.hasNext()) {
      holder.m_remainingHosts.add(hostIterator.next());
    }

    holder.m_component = pc.name;
    holder.m_tasks = Collections.singletonList(task);
    holder.m_preTasks = resolveTasks(upgradeContext, true, pc);
    holder.m_postTasks = resolveTasks(upgradeContext, false, pc);

    serviceToHostMap.put(service, holder);
  }

  @Override
  public List<StageWrapper> build(UpgradeContext upgradeContext, List<StageWrapper> stageWrappers) {

    if (0 == serviceToHostMap.size()) {
      return stageWrappers;
    }

    List<StageWrapper> starterUpgrades = new ArrayList<>();
    List<StageWrapper> finisherUpgrades = new ArrayList<>();

    // !!! create a stage wrapper for the first host tasks
    // !!! create a stage wrapper for the service check on the first host
    // !!! create a stage wrapper for the remaining hosts

    // !!! when building stages, the pre- and post- tasks should be run sequentially,
    // which means they should be in their own stages.  that may seem counter-intuitive,
    // but the parallelism is across hosts, not stages.

    serviceToHostMap.forEach((service, holder) -> {

      // !!! pre-tasks for the first host
      starterUpgrades.addAll(buildStageWrappers(upgradeContext, service, holder,
          holder.m_preTasks, true, "Preparing"));

      // !!! upgrades for the first host
      starterUpgrades.addAll(buildStageWrappers(upgradeContext, service, holder,
          holder.m_tasks, true, "Upgrading"));

      // !!! post tasks for the first host
      starterUpgrades.addAll(buildStageWrappers(upgradeContext, service, holder,
          holder.m_postTasks, true, "Completing"));

      // !!! service check for the first host
      StageWrapper serviceCheck = new ServiceCheckStageWrapper(service,
          upgradeContext.getServiceDisplay(service), false, holder.m_firstHost);
      starterUpgrades.add(serviceCheck);

      // !!! pre-tasks for remaining hosts
      finisherUpgrades.addAll(buildStageWrappers(upgradeContext, service, holder,
          holder.m_preTasks, false, "Prepare Remaining"));

      // !!! upgrades for the remaining hosts
      finisherUpgrades.addAll(buildStageWrappers(upgradeContext, service, holder,
          holder.m_tasks, false, "Upgrade Remaining"));

      // !!! post tasks for the remaining hosts
      finisherUpgrades.addAll(buildStageWrappers(upgradeContext, service, holder,
          holder.m_postTasks, false, "Complete Remaining"));
    });

    List<StageWrapper> results = new ArrayList<>(stageWrappers);

    results.addAll(starterUpgrades);
    results.addAll(finisherUpgrades);

    return results;
  }

  /**
   * Builds the stages for the components
   * @param upgradeContext
   *          the upgrade context
   * @param service
   *          the service name
   * @param holder
   *          the holder of component and hosts
   * @param tasks
   *          the tasks to wrap
   * @param firstHost
   *          {@code true} if wrapping for the first host
   * @param prefix
   *          the text prefix for the stage
   * @return
   *          the list of stage wrappers
   */
  private List<StageWrapper> buildStageWrappers(UpgradeContext upgradeContext, String service,
      HostHolder holder, List<Task> tasks, boolean firstHost, String prefix) {

    if (CollectionUtils.isEmpty(tasks)) {
      return Collections.emptyList();
    }

    Set<String> hosts = firstHost ? Collections.singleton(holder.m_firstHost) :
      holder.m_remainingHosts;

    String component = holder.m_component;
    String componentDisplay = upgradeContext.getComponentDisplay(service, component);
    String text = getStageText(prefix, componentDisplay, hosts);

    List<TaskWrapper> wrappers = buildTaskWrappers(service, component, tasks, hosts, firstHost);

    List<StageWrapper> results = new ArrayList<>();

    wrappers.forEach(task -> {
      // !!! there should only be one task per wrapper, so this assumption is ok for now
      StageWrapper.Type type = task.getTasks().get(0).getStageWrapperType();

      StageWrapper stage = new StageWrapper(type, text, new HashMap<>(),
          Collections.singletonList(task));
      results.add(stage);
    });

    return results;
  }

  /**
   * Build the wrappers for the tasks.
   *
   * @param service
   *          the service name
   * @param component
   *          the component name
   * @param tasks
   *          the tasks to wrap
   * @param hosts
   *          the hosts where the tasks should run
   * @param firstHost
   *          {@code true} if these wrappers are for the first host
   * @return
   *        the list if task wrappers
   */
  private List<TaskWrapper> buildTaskWrappers(String service, String component,
      List<Task> tasks, Set<String> hosts, boolean firstHost) {

    List<TaskWrapper> results = new ArrayList<>();

    tasks.forEach(task -> {
      // server side actions should run only run once (the first host)
      if (task.getType().isServerAction()) {
        if (firstHost) {
          String ambariServerHostname = StageUtils.getHostName();
          results.add(new TaskWrapper(service, component, Collections.singleton(ambariServerHostname), task));
        }
        return;
      }

      // FIXME how to handle master-only types

      // !!! the first host has already run tasks that are singular
      if (!firstHost && task.getType() == Type.EXECUTE) {
        ExecuteTask et = (ExecuteTask) task;

        // !!! singular types have already run when firstHost is true
        if (et.hosts != ExecuteHostType.ALL) {
          return;
        }
      }

      results.add(new TaskWrapper(service, component, hosts, task));
    });

    return results;
  }

  /**
   * Temporary holder for building stage wrappers
   */
  private static class HostHolder {
    private String m_component;
    private String m_firstHost;
    private Set<String> m_remainingHosts = new HashSet<>();
    // !!! there will only ever be one, but code is cleaner this way
    private List<Task> m_tasks;
    private List<Task> m_preTasks;
    private List<Task> m_postTasks;
  }

}
