/*
 * 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.tools.other;

import kafka.server.BrokerServer;
import kafka.server.KafkaBroker;
import kafka.utils.TestUtils;

import org.apache.kafka.clients.admin.Admin;
import org.apache.kafka.clients.admin.NewTopic;
import org.apache.kafka.clients.producer.KafkaProducer;
import org.apache.kafka.clients.producer.ProducerRecord;
import org.apache.kafka.common.MetricName;
import org.apache.kafka.common.Node;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.TopicPartitionInfo;
import org.apache.kafka.common.security.auth.SecurityProtocol;
import org.apache.kafka.common.serialization.ByteArraySerializer;
import org.apache.kafka.common.test.KafkaClusterTestKit;
import org.apache.kafka.common.test.TestKitNodes;
import org.apache.kafka.common.utils.Exit;
import org.apache.kafka.common.utils.Time;
import org.apache.kafka.common.utils.Utils;
import org.apache.kafka.image.TopicImage;
import org.apache.kafka.metadata.PartitionRegistration;
import org.apache.kafka.server.quota.QuotaType;
import org.apache.kafka.storage.internals.log.UnifiedLog;
import org.apache.kafka.tools.reassign.ReassignPartitionsCommand;

import org.apache.logging.log4j.core.config.Configurator;
import org.jfree.chart.ChartFactory;
import org.jfree.chart.ChartFrame;
import org.jfree.chart.JFreeChart;
import org.jfree.chart.plot.PlotOrientation;
import org.jfree.data.xy.XYSeries;
import org.jfree.data.xy.XYSeriesCollection;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.io.IOException;
import java.io.OutputStream;
import java.io.PrintWriter;
import java.nio.file.Files;
import java.text.DecimalFormat;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.TreeMap;
import java.util.concurrent.ExecutionException;
import java.util.function.Function;
import java.util.function.IntSupplier;
import java.util.stream.Collectors;
import java.util.stream.IntStream;

import javax.imageio.ImageIO;

import scala.Option;

import static java.nio.file.StandardOpenOption.APPEND;
import static java.nio.file.StandardOpenOption.CREATE;
import static org.apache.kafka.test.TestUtils.DEFAULT_MAX_WAIT_MS;

/**
 * Test rig for measuring throttling performance. Configure the parameters for a set of experiments, then execute them
 * and view the html output file, with charts, that are produced. You can also render the charts to the screen if
 * you wish.
 * <p>
 * Currently you'll need about 40GB of disk space to run these experiments (largest data written x2). Tune the msgSize
 * & #partitions and throttle to adjust this.
 */
public class ReplicationQuotasTestRig {
    public static final Logger LOGGER = LoggerFactory.getLogger(ReplicationQuotasTestRig.class);

    public static final int K = 1000 * 1000;

    private static final String DIR;

    static {
        Configurator.reconfigure();

        new File("Experiments").mkdir();
        DIR = "Experiments/Run" + Long.valueOf(System.currentTimeMillis()).toString().substring(8);
        new File(DIR).mkdir();
    }

    public static void main(String[] args) {
        boolean displayChartsOnScreen = args.length > 0 && Objects.equals(args[0], "show-gui");
        Journal journal = new Journal();

        List<ExperimentDef> experiments = List.of(
            //1GB total data written, will take 210s
            new ExperimentDef("Experiment1", 5, 20, 1 * K, 500, 100 * 1000),
            //5GB total data written, will take 110s
            new ExperimentDef("Experiment2", 5, 50, 10 * K, 1000, 100 * 1000),
            //5GB total data written, will take 110s
            new ExperimentDef("Experiment3", 50, 50, 2 * K, 1000, 100 * 1000),
            //10GB total data written, will take 110s
            new ExperimentDef("Experiment4", 25, 100, 4 * K, 1000, 100 * 1000),
            //10GB total data written, will take 80s
            new ExperimentDef("Experiment5", 5, 50, 50 * K, 4000, 100 * 1000)
        );
        experiments.forEach(def -> run(def, journal, displayChartsOnScreen));

        if (!displayChartsOnScreen)
            Exit.exit(0);
    }

    static void run(ExperimentDef config, Journal journal, boolean displayChartsOnScreen) {
        Experiment experiment = new Experiment();
        try {
            experiment.run(config, journal, displayChartsOnScreen);
            journal.footer();
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            experiment.tearDown();
        }
    }

    static class ExperimentDef {
        String name;
        int brokers;
        int partitions;
        long throttle;
        int msgsPerPartition;
        int msgSize;
        final long targetBytesPerBrokerMB;

        public ExperimentDef(String name, int brokers, int partitions, long throttle, int msgsPerPartition, int msgSize) {
            this.name = name;
            this.brokers = brokers;
            this.partitions = partitions;
            this.throttle = throttle;
            this.msgsPerPartition = msgsPerPartition;
            this.msgSize = msgSize;
            this.targetBytesPerBrokerMB = (long) msgsPerPartition * (long) msgSize * (long) partitions / brokers / 1_000_000;
        }
    }

    static class Experiment {
        static final String TOPIC_NAME = "my-topic";

        String experimentName = "unset";
        Map<Integer, List<Double>> leaderRates = new HashMap<>();
        Map<Integer, List<Double>> followerRates = new HashMap<>();
        KafkaClusterTestKit cluster;
        Admin adminClient;

        void startBrokers(int numBrokerNodes) {
            System.out.println("Starting Brokers");

            try {
                cluster = new KafkaClusterTestKit.Builder(
                        new TestKitNodes.Builder()
                                .setNumControllerNodes(1)
                                .setNumBrokerNodes(numBrokerNodes)
                                .build()
                ).build();
                cluster.format();
                cluster.startup();
                cluster.waitForReadyBrokers();
            } catch (Exception e) {
                throw new RuntimeException("Failed to start test Kafka cluster", e);
            }

            adminClient = Admin.create(cluster.clientProperties());
        }

        public void tearDown() {
            Utils.closeQuietly(adminClient, "adminClient");
            try {
                cluster.close();
            } catch (Exception e) {
                throw new RuntimeException(e);
            }
        }

        public void run(ExperimentDef config, Journal journal, boolean displayChartsOnScreen) throws Exception {
            experimentName = config.name;
            int shift = Math.round(config.brokers / 2f);

            IntSupplier nextReplicaRoundRobin = new IntSupplier() {
                int count = 0;

                @Override public int getAsInt() {
                    count++;
                    return (count + shift) % config.brokers;
                }
            };

            Map<Integer, List<Integer>> replicas = IntStream.rangeClosed(0, config.partitions - 1).boxed().collect(Collectors.toMap(
                Function.identity(),
                partition -> List.of(nextReplicaRoundRobin.getAsInt())
            ));

            startBrokers(config.brokers);
            adminClient.createTopics(Set.of(new NewTopic(TOPIC_NAME, replicas))).all().get();

            TestUtils.waitUntilTrue(
                    () -> cluster.brokers().values().stream().allMatch(server -> {
                        TopicImage image = server.metadataCache().currentImage().topics().getTopic(TOPIC_NAME);
                        return image != null && image.partitions().values().stream().allMatch(PartitionRegistration::hasLeader);
                    }),
                    () -> "Timed out waiting for topic listing",
                    DEFAULT_MAX_WAIT_MS,
                    500L
            );

            System.out.println("Writing Data");
            try (KafkaProducer<byte[], byte[]> producer = createProducer()) {
                for (int x = 0; x < config.msgsPerPartition; x++) {
                    for (int partition = 0; partition < config.partitions; partition++) {
                        producer.send(new ProducerRecord<>(TOPIC_NAME, partition, null, new byte[config.msgSize]));
                    }
                }
            }

            System.out.println("Generating Reassignment");
            Map<TopicPartition, List<Integer>> newAssignment = ReassignPartitionsCommand.generateAssignment(
                    adminClient,
                    json(TOPIC_NAME),
                    cluster.brokers().values().stream()
                            .map(server -> String.valueOf(server.replicaManager().localBrokerId()))
                            .collect(Collectors.joining(",")),
                    true
            ).getKey();

            System.out.println("Starting Reassignment");
            long start = System.currentTimeMillis();

            ReassignPartitionsCommand.executeAssignment(adminClient, false,
                ReassignPartitionsCommand.formatAsReassignmentJson(newAssignment, Map.of()),
                config.throttle, -1L, 10000L, Time.SYSTEM, false);

            //Await completion
            waitForReassignmentToComplete();
            System.out.println("Reassignment took " + (System.currentTimeMillis() - start) / 1000 + "s");

            validateAllOffsetsMatch(config);

            journal.appendToJournal(config);
            renderChart(leaderRates, "Leader", journal, displayChartsOnScreen);
            renderChart(followerRates, "Follower", journal, displayChartsOnScreen);
            logOutput(config, replicas, newAssignment);

            System.out.println("Output can be found here: " + journal.path());
        }

        void validateAllOffsetsMatch(ExperimentDef config) {
            //Validate that offsets are correct in all brokers
            for (KafkaBroker broker : cluster.brokers().values()) {
                for (int partitionId = 0; partitionId < config.partitions; partitionId++) {
                    long offset = broker.logManager().getLog(new TopicPartition(TOPIC_NAME, partitionId), false)
                            .map(UnifiedLog::logEndOffset).getOrElse(() -> -1L);
                    if (offset >= 0 && offset != config.msgsPerPartition) {
                        throw new RuntimeException(
                            "Run failed as offsets did not match for partition " + partitionId + " on broker " + broker.config().nodeId() + ". " +
                            "Expected " + config.msgsPerPartition + " but was " + offset + "."
                        );
                    }
                }
            }
        }

        void logOutput(ExperimentDef config, Map<Integer, List<Integer>> replicas, Map<TopicPartition, List<Integer>> newAssignment) throws Exception {
            List<TopicPartitionInfo> actual = adminClient.describeTopics(Set.of(TOPIC_NAME))
                .allTopicNames().get().get(TOPIC_NAME).partitions();

            Map<Integer, List<Integer>> curAssignment = actual.stream().collect(Collectors.toMap(
                TopicPartitionInfo::partition,
                p -> p.replicas().stream().map(Node::id).toList()
            ));

            //Long stats
            System.out.println("The replicas are " + new TreeMap<>(replicas).entrySet().stream().map(e -> "\n" + e).collect(Collectors.joining()));
            System.out.println("This is the current replica assignment:\n" + curAssignment);
            System.out.println("proposed assignment is: \n" + newAssignment);
            System.out.println("This is the assignment we ended up with " + curAssignment);

            //Test Stats
            System.out.println("numBrokers: " + config.brokers);
            System.out.println("numPartitions: " + config.partitions);
            System.out.println("throttle: " + config.throttle);
            System.out.println("numMessagesPerPartition: " + config.msgsPerPartition);
            System.out.println("msgSize: " + config.msgSize);
            System.out.println("We will write " + config.targetBytesPerBrokerMB + "MB of data per broker");
            System.out.println("Worst case duration is " + config.targetBytesPerBrokerMB * 1000 * 1000 / config.throttle);
        }

        void waitForReassignmentToComplete() {
            TestUtils.waitUntilTrue(() -> {
                printRateMetrics();
                try {
                    return adminClient.listPartitionReassignments().reassignments().get().isEmpty();
                } catch (InterruptedException | ExecutionException e) {
                    throw new RuntimeException(e);
                }
            }, () -> "Partition reassignments didn't complete.", 60 * 60 * 1000, 1000L);
        }

        void renderChart(Map<Integer, List<Double>> data, String name, Journal journal, boolean displayChartsOnScreen) throws Exception {
            XYSeriesCollection dataset = addDataToChart(data);
            JFreeChart chart = createChart(name, dataset);

            writeToFile(name, journal, chart);
            maybeDisplayOnScreen(displayChartsOnScreen, chart);
            System.out.println("Chart generated for " + name);
        }

        void maybeDisplayOnScreen(boolean displayChartsOnScreen, JFreeChart chart) {
            if (displayChartsOnScreen) {
                ChartFrame frame = new ChartFrame(experimentName, chart);
                frame.pack();
                frame.setVisible(true);
            }
        }

        void writeToFile(String name, Journal journal, JFreeChart chart) throws Exception {
            File file = new File(DIR, experimentName + "-" + name + ".png");
            ImageIO.write(chart.createBufferedImage(1000, 700), "png", file);
            journal.appendChart(file.getAbsolutePath(), name.equals("Leader"));
        }

        JFreeChart createChart(String name, XYSeriesCollection dataset) {
            return ChartFactory.createXYLineChart(
                experimentName + " - " + name + " Throttling Performance",
                "Time (s)",
                "Throttle Throughput (B/s)",
                dataset, PlotOrientation.VERTICAL, false, true, false
            );
        }

        XYSeriesCollection addDataToChart(Map<Integer, List<Double>> data) {
            XYSeriesCollection dataset = new XYSeriesCollection();
            data.forEach((broker, values) -> {
                XYSeries series = new XYSeries("Broker:" + broker);
                int x = 0;
                for (double value : values) {
                    series.add(x, value);
                    x++;
                }
                dataset.addSeries(series);
            });
            return dataset;
        }

        void record(Map<Integer, List<Double>> rates, int nodeId, Double currentRate) {
            List<Double> leaderRatesBroker = rates.getOrDefault(nodeId, new ArrayList<>());
            leaderRatesBroker.add(currentRate);
            rates.put(nodeId, leaderRatesBroker);
        }

        void printRateMetrics() {
            for (BrokerServer broker : cluster.brokers().values()) {
                double leaderRate = measuredRate(broker, QuotaType.LEADER_REPLICATION);
                if (broker.config().nodeId() == 0)
                    LOGGER.info("waiting... Leader rate on 1 is {}", leaderRate);
                record(leaderRates, broker.config().nodeId(), leaderRate);
                if (leaderRate > 0)
                    LOGGER.trace("Leader Rate on {} is {}", broker.config().nodeId(), leaderRate);

                double followerRate = measuredRate(broker, QuotaType.FOLLOWER_REPLICATION);
                record(followerRates, broker.config().nodeId(), followerRate);
                if (followerRate > 0)
                    LOGGER.trace("Follower Rate on {} is {}", broker.config().nodeId(), followerRate);
            }
        }

        private double measuredRate(KafkaBroker broker, QuotaType repType) {
            MetricName metricName = broker.metrics().metricName("byte-rate", repType.toString());
            return broker.metrics().metrics().containsKey(metricName)
                ? (double) broker.metrics().metrics().get(metricName).metricValue()
                : -1d;
        }

        String json(String... topic) {
            String topicStr = Arrays.stream(topic).map(t -> "{\"topic\": \"" + t + "\"}").collect(Collectors.joining(","));
            return "{\"topics\": [" + topicStr + "],\"version\":1}";
        }

        KafkaProducer<byte[], byte[]> createProducer() {
            return TestUtils.createProducer(
                cluster.bootstrapServers(),
                1,
                60 * 1000L,
                1024L * 1024L,
                Integer.MAX_VALUE,
                30 * 1000,
                0,
                16384,
                "none",
                20 * 1000,
                SecurityProtocol.PLAINTEXT,
                Option.empty(),
                Option.empty(),
                new ByteArraySerializer(),
                new ByteArraySerializer(),
                false
            );
        }
    }

    static class Journal {
        File log = new File(DIR, "Log.html");

        public Journal() {
            header();
        }

        void appendToJournal(ExperimentDef config) {
            DecimalFormat format = new DecimalFormat("###,###.###");

            String message = "\n\n<h3>" + config.name + "</h3>" +
                "<p>- BrokerCount: " + config.brokers +
                "<p>- PartitionCount: " + config.partitions +
                "<p>- Throttle: " + format.format(config.throttle) + " MB/s" +
                "<p>- MsgCount: " + format.format(config.msgsPerPartition) + " " +
                "<p>- MsgSize: " + format.format(config.msgSize) +
                "<p>- TargetBytesPerBrokerMB: " + config.targetBytesPerBrokerMB + "<p>";

            append(message);
        }

        void appendChart(String path, boolean first) {
            StringBuilder message = new StringBuilder();
            if (first)
                message.append("<p><p>");
            message.append("<img src=\"" + path + "\" alt=\"Chart\" style=\"width:600px;height:400px;align=\"middle\"\">");
            if (!first)
                message.append("<p><p>");
            append(message.toString());
        }

        void header() {
            append("<html><head><h1>Replication Quotas Test Rig</h1></head><body>");
        }

        void footer() {
            append("</body></html>");
        }

        void append(String message) {
            try {
                OutputStream stream = Files.newOutputStream(log.toPath(), CREATE, APPEND);
                PrintWriter writer = new PrintWriter(stream);
                writer.append(message);
                writer.close();
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
        }

        String path() {
            return log.getAbsolutePath();
        }
    }
}
