package resa.evaluation.scheduler.plan;

import backtype.storm.Config;
import backtype.storm.generated.Nimbus;
import backtype.storm.generated.RebalanceOptions;
import backtype.storm.scheduler.ExecutorDetails;
import backtype.storm.utils.NimbusClient;
import backtype.storm.utils.Utils;
import org.apache.curator.framework.CuratorFramework;
import org.apache.zookeeper.ZKUtil;
import org.junit.Before;
import org.junit.Test;
import resa.evaluation.migrate.ConsistentHashing;
import resa.evaluation.migrate.MigPlan;
import resa.util.TopologyHelper;

import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.*;
import java.util.stream.Collectors;

/**
 * Created by ding on 14-8-12.
 */
public class MigrateSimulateForCH {

    private CuratorFramework zk;
    private Map<String, Object> conf;
    private int[] allocations;
    private List<String> slots;
    private double[] dataSizes;
    private ExecutorDetails[] executors;
    private String topoName = "fpt";

    @Before
    public void init() throws Exception {
        conf = Utils.readStormConfig();
        conf.put(Config.NIMBUS_HOST, "192.168.0.30");
        conf.put(Config.STORM_ZOOKEEPER_SERVERS, Arrays.asList("192.168.0.30"));
        zk = Utils.newCuratorStarted(conf, (List<String>) conf.get(Config.STORM_ZOOKEEPER_SERVERS),
                conf.get(Config.STORM_ZOOKEEPER_PORT));
        allocations = Files.readAllLines(Paths.get("/Volumes/Data/work/doctor/resa/exp/ch-migrate.txt")).stream()
                .map(String::trim).filter(s -> !s.isEmpty())
                .mapToInt(Integer::parseInt).toArray();
        dataSizes = Files.readAllLines(Paths.get("/Volumes/Data/work/doctor/resa/exp/fp-data-sizes-064-70.txt")).stream()
                .map(String::trim).filter(s -> !s.isEmpty()).mapToDouble(Double::valueOf).toArray();
        slots = getSlots();
        String topoId = TopologyHelper.getTopologyId(topoName, conf);
        Map<String, Map<ExecutorDetails, String>> assignment = TopologyHelper.getAssignment(conf, topoId);
        Map<ExecutorDetails, String> compAssignment = assignment.get("detector");
        executors = compAssignment.keySet().stream().sorted(Comparator.comparing(ExecutorDetails::getStartTask))
                .toArray(ExecutorDetails[]::new);
    }

    @Test
    public void deleteZkNode() throws Exception {
        ZKUtil.deleteRecursive(zk.getZookeeperClient().getZooKeeper(), "/resa");
    }

    @Test
    public void run() throws Exception {
        int state = 0;
        List<int[]> from = new ArrayList<>();
        List<int[]> to = new ArrayList<>();
        String[] workers = slots.stream().limit(allocations[state]).toArray(String[]::new);
        assignment2State(from, workers);
        doRebalance(topoName, allocations[state]);
        System.out.println("Init state, curr alloc=" + allocations[state]);
        Scanner scanner = new Scanner(System.in);
        String line;
        int count = 0;
        while ((line = scanner.nextLine()) != null) {
            if (line.equals("c")) {
                count++;
                System.out.println("run " + count + ", from " + allocations[state] + " to " + allocations[state + 1]);
                System.out.println("old workers:" + Arrays.toString(workers));
                workers = migrateTo(from, workers, to, allocations[state + 1]);
                System.out.println("new workers:" + Arrays.toString(workers));
                assignment2State(to, workers);
                doRebalance(topoName, allocations[state + 1]);
                state++;
                from.clear();
                from.addAll(to);
                to.clear();
            }
        }
    }

    private String[] migrateTo(List<int[]> from, String[] workers, List<int[]> to, int toNumWorkers) {
        MigPlan res = ConsistentHashing.calc(dataSizes, from, toNumWorkers);
        to.addAll(res.tasks);
        String[] newWorkers = new String[toNumWorkers];
        Set<String> selected = new HashSet<>();
        int[][] map = res.matching;
        for (int i = 0; i < map.length; i++) {
            newWorkers[map[i][0]] = workers[map[i][1]];
            selected.add(newWorkers[map[i][0]]);
        }
        for (int i = 0; i < newWorkers.length; i++) {
            if (newWorkers[i] == null) {
                for (String w : slots) {
                    if (!selected.contains(w)) {
                        selected.add(w);
                        newWorkers[i] = w;
                        break;
                    }
                }
            }
        }
        return newWorkers;
    }

    public void assignment2State(List<int[]> alloc, String[] workers) throws Exception {
        Map<String, String> exe2WorkerSlot = new HashMap<>();
        for (int i = 0; i < alloc.size(); i++) {
            for (int task : alloc.get(i)) {
                exe2WorkerSlot.put(executors[task].getStartTask() + "-" + executors[task].getEndTask(), workers[i]);
            }
        }
//        System.out.println(exe2WorkerSlot);
        String path = "/resa/" + TopologyHelper.getTopologyId(topoName, conf);
        if (zk.checkExists().forPath(path) != null) {
            zk.setData().forPath(path, Utils.serialize(new HashMap<>(Collections.singletonMap("assignment",
                    exe2WorkerSlot))));
        } else {
            zk.create().creatingParentsIfNeeded().forPath(path, Utils.serialize(
                    new HashMap<>(Collections.singletonMap("assignment", exe2WorkerSlot))));
        }
    }


    private List<String> getSlots() throws Exception {
        NimbusClient nimbusClient = NimbusClient.getConfiguredClient(conf);
        Nimbus.Client nimbus = nimbusClient.getClient();
        List<String> slots = nimbus.getClusterInfo().get_supervisors().stream()
                .flatMap(s -> Arrays.asList(s.get_host() + ":6700", s.get_host() + ":6701").stream())
                .collect(Collectors.toList());
        Collections.shuffle(slots);
        System.out.println(slots.size());
        return slots;
    }

    private void doRebalance(String topoName, int numWorkers) throws Exception {
        NimbusClient nimbusClient = NimbusClient.getConfiguredClient(conf);
        Nimbus.Client nimbus = nimbusClient.getClient();
        RebalanceOptions options = new RebalanceOptions();
        options.set_num_workers(numWorkers);
        options.set_wait_secs(1);
        System.out.println("Reassigning to " + numWorkers + " workers");
        nimbus.rebalance(topoName, options);
    }

    public static void main(String[] args) throws Exception {
        MigrateSimulateForCH simulate = new MigrateSimulateForCH();
        simulate.init();
        simulate.run();
    }

}
