/*
 * 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.kafka.trogdor.fault;

import com.fasterxml.jackson.databind.node.TextNode;
import org.apache.kafka.common.internals.KafkaFutureImpl;
import org.apache.kafka.trogdor.common.Node;
import org.apache.kafka.trogdor.common.Platform;
import org.apache.kafka.trogdor.task.TaskWorker;
import org.apache.kafka.trogdor.task.WorkerStatusTracker;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.net.NetworkInterface;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.function.Consumer;
import java.util.stream.Stream;

/**
 * Uses the linux utility <pre>tc</pre> (traffic controller) to degrade performance on a specified network device
 */
public class DegradedNetworkFaultWorker implements TaskWorker {

    private static final Logger log = LoggerFactory.getLogger(DegradedNetworkFaultWorker.class);

    private final String id;
    private final Map<String, DegradedNetworkFaultSpec.NodeDegradeSpec> nodeSpecs;
    private WorkerStatusTracker status;

    public DegradedNetworkFaultWorker(String id, Map<String, DegradedNetworkFaultSpec.NodeDegradeSpec> nodeSpecs) {
        this.id = id;
        this.nodeSpecs = nodeSpecs;
    }

    @Override
    public void start(Platform platform, WorkerStatusTracker status, KafkaFutureImpl<String> haltFuture) throws Exception {
        log.info("Activating DegradedNetworkFaultWorker {}.", id);
        this.status = status;
        this.status.update(new TextNode("enabling traffic control " + id));
        Node curNode = platform.curNode();
        DegradedNetworkFaultSpec.NodeDegradeSpec nodeSpec = nodeSpecs.get(curNode.name());
        if (nodeSpec != null) {
            for (String device : devicesForSpec(nodeSpec)) {
                if (nodeSpec.latencyMs() < 0 || nodeSpec.rateLimitKbit() < 0) {
                    throw new RuntimeException("Expected non-negative values for latencyMs and rateLimitKbit, but got " + nodeSpec);
                } else {
                    enableTrafficControl(platform, device, nodeSpec.latencyMs(), nodeSpec.rateLimitKbit());
                }
            }
        }
        this.status.update(new TextNode("enabled traffic control " + id));
    }

    @Override
    public void stop(Platform platform) throws Exception {
        log.info("Deactivating DegradedNetworkFaultWorker {}.", id);
        this.status.update(new TextNode("disabling traffic control " + id));
        Node curNode = platform.curNode();
        DegradedNetworkFaultSpec.NodeDegradeSpec nodeSpec = nodeSpecs.get(curNode.name());
        if (nodeSpec != null) {
            for (String device : devicesForSpec(nodeSpec)) {
                disableTrafficControl(platform, device);
            }
        }
        this.status.update(new TextNode("disabled traffic control " + id));
    }

    private Set<String> devicesForSpec(DegradedNetworkFaultSpec.NodeDegradeSpec nodeSpec) throws Exception {
        Set<String> devices = new HashSet<>();
        if (nodeSpec.networkDevice().isEmpty()) {
            for (NetworkInterface networkInterface : Collections.list(NetworkInterface.getNetworkInterfaces())) {
                if (!networkInterface.isLoopback()) {
                    devices.add(networkInterface.getName());
                }
            }
        } else {
            devices.add(nodeSpec.networkDevice());
        }
        return devices;
    }

    /**
     * Constructs the appropriate "tc" commands to apply latency and rate limiting, if they are non zero.
     */
    private void enableTrafficControl(Platform platform, String networkDevice, int delayMs, int rateLimitKbps) throws IOException {
        if (delayMs > 0) {
            int deviationMs = Math.max(1, (int) Math.sqrt(delayMs));
            List<String> delay = new ArrayList<>();
            rootHandler(networkDevice, delay::add);
            netemDelay(delayMs, deviationMs, delay::add);
            platform.runCommand(delay.toArray(new String[]{}));

            if (rateLimitKbps > 0) {
                List<String> rate = new ArrayList<>();
                childHandler(networkDevice, rate::add);
                tbfRate(rateLimitKbps, rate::add);
                platform.runCommand(rate.toArray(new String[]{}));
            }
        } else if (rateLimitKbps > 0) {
            List<String> rate = new ArrayList<>();
            rootHandler(networkDevice, rate::add);
            tbfRate(rateLimitKbps, rate::add);
            platform.runCommand(rate.toArray(new String[]{}));
        } else {
            log.warn("Not applying any rate limiting or latency");
        }
    }

    /**
     * Construct the first part of a "tc" command to define a qdisc root handler for the given network interface
     */
    private void rootHandler(String networkDevice, Consumer<String> consumer) {
        Stream.of("sudo", "tc", "qdisc", "add", "dev", networkDevice, "root", "handle", "1:0").forEach(consumer);
    }

    /**
     * Construct the first part of a "tc" command to define a qdisc child handler for the given interface. This can
     * only be used if a root handler has been appropriately defined first (as in {@link #rootHandler}).
     */
    private void childHandler(String networkDevice, Consumer<String> consumer) {
        Stream.of("sudo", "tc", "qdisc", "add", "dev", networkDevice, "parent", "1:1", "handle", "10:").forEach(consumer);
    }

    /**
     * Construct the second part of a "tc" command that defines a netem (Network Emulator) filter that will apply some
     * amount of latency with a small amount of deviation. The distribution of the latency deviation follows a so-called
     * Pareto-normal distribution. This is the formal name for the 80/20 rule, which might better represent real-world
     * patterns.
     */
    private void netemDelay(int delayMs, int deviationMs, Consumer<String> consumer) {
        Stream.of("netem", "delay", String.format("%dms", delayMs), String.format("%dms", deviationMs),
                "distribution", "paretonormal").forEach(consumer);
    }

    /**
     * Construct the second part of a "tc" command that defines a tbf (token buffer filter) that will rate limit the
     * packets going through a qdisc.
     */
    private void tbfRate(int rateLimitKbit, Consumer<String> consumer) {
        Stream.of("tbf", "rate", String.format("%dkbit", rateLimitKbit), "burst", "1mbit", "latency", "500ms").forEach(consumer);
    }

    /**
     * Delete any previously defined qdisc for the given network interface.
     * @throws IOException
     */
    private void disableTrafficControl(Platform platform, String networkDevice) throws IOException {
        platform.runCommand(new String[] {
            "sudo", "tc", "qdisc", "del", "dev", networkDevice, "root"
        });
    }
}
