// Copyright 2016 Twitter. All rights reserved.
//
// 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 com.twitter.heron.scheduler.aurora;

import java.net.HttpURLConnection;
import java.util.Map;
import java.util.Set;
import java.util.logging.Logger;

import com.twitter.heron.proto.system.PhysicalPlans.StMgr;
import com.twitter.heron.spi.common.Config;
import com.twitter.heron.spi.common.Context;
import com.twitter.heron.spi.packing.PackingPlan;
import com.twitter.heron.spi.statemgr.IStateManager;
import com.twitter.heron.spi.statemgr.SchedulerStateManagerAdaptor;
import com.twitter.heron.spi.utils.NetworkUtils;
import com.twitter.heron.spi.utils.ReflectionUtils;

/**
 * Implementation of AuroraController that is a wrapper of AuroraCLIController.
 * The difference is `restart` command:
 * 1. restart whole topology: delegate to AuroraCLIController
 * 2. restart container 0: delegate to AuroraCLIController
 * 3. restart container x(x>0): call heron-shell endpoint `/killexecutor`
 * For backpressure, only containers with heron-stmgr may send out backpressure.
 * This class is to handle `restart backpressure containers inside container`,
 * while delegating to AuroraCLIController for all the other scenarios.
 */
class AuroraHeronShellController implements AuroraController {
  private static final Logger LOG = Logger.getLogger(AuroraHeronShellController.class.getName());

  private final String topologyName;
  private final AuroraCLIController cliController;
  private final SchedulerStateManagerAdaptor stateMgrAdaptor;

  AuroraHeronShellController(String jobName, String cluster, String role, String env,
      String auroraFilename, boolean isVerbose, Config localConfig)
      throws ClassNotFoundException, InstantiationException, IllegalAccessException {
    this.topologyName = jobName;
    this.cliController =
        new AuroraCLIController(jobName, cluster, role, env, auroraFilename, isVerbose);

    Config config = Config.toClusterMode(localConfig);
    String stateMgrClass = Context.stateManagerClass(config);
    IStateManager stateMgr = ReflectionUtils.newInstance(stateMgrClass);
    stateMgr.initialize(config);
    stateMgrAdaptor = new SchedulerStateManagerAdaptor(stateMgr, 5000);
  }

  @Override
  public boolean createJob(Map<AuroraField, String> bindings) {
    return cliController.createJob(bindings);
  }

  @Override
  public boolean killJob() {
    return cliController.killJob();
  }

  // Restart an aurora container
  @Override
  public boolean restart(Integer containerId) {
    // there is no backpressure for container 0, delegate to aurora client
    if (containerId == null || containerId == 0) {
      cliController.restart(containerId);
    }

    if (stateMgrAdaptor == null) {
      LOG.warning("SchedulerStateManagerAdaptor not initialized");
      return false;
    }

    int index = containerId - 1; // stmgr container starts from 1
    StMgr contaienrInfo = stateMgrAdaptor.getPhysicalPlan(topologyName).getStmgrs(index);
    String host = contaienrInfo.getHostName();
    int port = contaienrInfo.getShellPort();
    String url = "http://" + host + ":" + port + "/killexecutor";
    String payload = "secret=" + stateMgrAdaptor.getExecutionState(topologyName).getTopologyId();
    LOG.info("sending `kill container` to " + url + "; payload: " + payload);

    HttpURLConnection con = NetworkUtils.getHttpConnection(url);
    try {
      NetworkUtils.sendHttpPostRequest(con, "X", payload.getBytes());
      return NetworkUtils.checkHttpResponseCode(con, 200);
    } finally {
      con.disconnect();
    }
  }

  @Override
  public void removeContainers(Set<PackingPlan.ContainerPlan> containersToRemove) {
    cliController.removeContainers(containersToRemove);
  }

  @Override
  public void addContainers(Integer count) {
    cliController.addContainers(count);
  }
}

@@ -2940,16 +2940,276 @@
 );%0A  %7D%0A%0A
+  private StMgr searchContainer(Integer id) %7B%0A    String prefix = %22stmgr-%22 + id;%0A    for (StMgr sm : stateMgrAdaptor.getPhysicalPlan(topologyName).getStmgrsList()) %7B%0A      if (sm.getId().equals(prefix)) %7B%0A        return sm;%0A      %7D%0A    %7D%0A    return null;%0A  %7D%0A%0A
   // Res
@@ -3421,24 +3421,31 @@
  0) %7B%0A      
+return 
 cliControlle
@@ -3612,298 +3612,235 @@
     
-int index = containerId - 1; // stmgr container starts from 1%0A    StMgr contaienrInfo = stateMgrAdaptor.getPhysicalPlan(topologyName).getStmgrs(index);%0A    String host = contaienrInfo.getHostName();%0A    int port = contaienrInfo.getShellPort();%0A    String url = %22http://%22 + host + %22:%22 + p
+StMgr sm = searchContainer(containerId);%0A    if (sm == null) %7B%0A      LOG.warning(%22container not found in pplan %22 + containerId);%0A      return false;%0A    %7D%0A%0A    String url = %22http://%22 + sm.getHostName() + %22:%22 + sm.getShellP
 ort
+()
  + %22
@@ -4113,16 +4113,20 @@
 %7B%0A      
+if (
 NetworkU
@@ -4179,18 +4179,22 @@
 Bytes())
-;%0A
+) %7B%0A  
       re
@@ -4244,16 +4244,224 @@
 , 200);%0A
+      %7D else %7B // if heron-shell command fails, delegate to aurora client%0A        LOG.info(%22heron-shell killexecutor failed; try aurora client ..%22);%0A        return cliController.restart(containerId);%0A      %7D%0A
     %7D fi

fix bugs in AuroraHeronShellController (#2553)