/*
 * 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.clients.admin;

import org.apache.kafka.clients.producer.Producer;
import org.apache.kafka.clients.producer.ProducerRecord;
import org.apache.kafka.common.Node;
import org.apache.kafka.common.TopicPartitionInfo;
import org.apache.kafka.common.test.ClusterInstance;
import org.apache.kafka.common.test.api.ClusterConfigProperty;
import org.apache.kafka.common.test.api.ClusterTest;
import org.apache.kafka.common.test.api.ClusterTestDefaults;
import org.apache.kafka.common.test.api.Type;
import org.apache.kafka.server.config.ReplicationConfigs;
import org.apache.kafka.server.config.ServerConfigs;
import org.apache.kafka.server.config.ServerLogConfigs;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

import static org.junit.jupiter.api.Assertions.assertEquals;

@ClusterTestDefaults(types = {Type.KRAFT},
    brokers = 4,
    serverProperties = {
        @ClusterConfigProperty(key = ServerLogConfigs.NUM_PARTITIONS_CONFIG, value = "8"),
        @ClusterConfigProperty(key = ReplicationConfigs.DEFAULT_REPLICATION_FACTOR_CONFIG, value = "2"),
        @ClusterConfigProperty(id = 0, key = ServerConfigs.BROKER_RACK_CONFIG, value = "0"),
        @ClusterConfigProperty(id = 1, key = ServerConfigs.BROKER_RACK_CONFIG, value = "0"),
        @ClusterConfigProperty(id = 2, key = ServerConfigs.BROKER_RACK_CONFIG, value = "1"),
        @ClusterConfigProperty(id = 3, key = ServerConfigs.BROKER_RACK_CONFIG, value = "1"),
    })
public class RackAwareAutoTopicCreationTest {

    private static final String TOPIC = "topic";

    @ClusterTest
    public void testAutoCreateTopic(ClusterInstance cluster) throws Exception {

        try (Admin admin = cluster.admin();
             Producer<byte[], byte[]> producer = cluster.producer()) {

            // send a record to trigger auto create topic
            ProducerRecord<byte[], byte[]> record = new ProducerRecord<>(TOPIC, null, "key".getBytes(), "value".getBytes());
            assertEquals(0L, producer.send(record).get().offset(), "Should have offset 0");

            // check broker rack content
            Map<Integer, String> expectedBrokerToRackMap = Map.of(
                    0, "0",
                    1, "0",
                    2, "1",
                    3, "1"
            );
            Map<Integer, String> actualBrokerToRackMap = getBrokerToRackMap(cluster);
            assertEquals(expectedBrokerToRackMap, actualBrokerToRackMap);

            // get topic assignments and check it's content
            Map<Integer, List<Integer>> assignments = getTopicAssignment(admin);
            for (List<Integer> brokerList : assignments.values()) {
                assertEquals(new HashSet<>(brokerList).size(), brokerList.size(),
                        "More than one replica is assigned to same broker for the same partition");
            }

            // check rack count for each partition
            ReplicaDistributions distribution = getReplicaDistribution(assignments, expectedBrokerToRackMap);
            Map<String, String> serverProperties = cluster.config().serverProperties();
            int numPartition = Integer.parseInt(serverProperties.get(ServerLogConfigs.NUM_PARTITIONS_CONFIG));
            int replicationFactor = Integer.parseInt(serverProperties.get(ReplicationConfigs.DEFAULT_REPLICATION_FACTOR_CONFIG));
            List<Integer> expectedRackCounts = Collections.nCopies(numPartition, replicationFactor);
            List<Integer> actualRackCounts = distribution.partitionToRackMap.values().stream()
                    .map(racks -> (int) racks.stream().distinct().count())
                    .collect(Collectors.toList());
            assertEquals(expectedRackCounts, actualRackCounts,
                    "More than one replica of the same partition is assigned to the same rack");

            // check replica count for each partition
            int numBrokers = cluster.brokers().size();
            int numReplicasPerBroker = numPartition * replicationFactor / numBrokers;
            List<Integer> expectedReplicasCounts = Collections.nCopies(numBrokers, numReplicasPerBroker);
            List<Integer> actualReplicasCounts = new ArrayList<>(distribution.partitionToCountMap.values());
            assertEquals(expectedReplicasCounts, actualReplicasCounts, "Replica count is not even for broker");
        }
    }

    private static Map<Integer, List<Integer>> getTopicAssignment(Admin admin) throws Exception {
        TopicDescription topicDescription = admin.describeTopics(List.of(TOPIC)).allTopicNames().get().get(TOPIC);
        return topicDescription.partitions().stream()
                .collect(Collectors.toMap(
                        TopicPartitionInfo::partition,
                        p -> p.replicas().stream().map(Node::id).collect(Collectors.toList())));
    }

    private static Map<Integer, String> getBrokerToRackMap(ClusterInstance cluster) throws Exception {
        Map<Integer, String> actualBrokerToRackMap = new HashMap<>();
        try (Admin admin = cluster.admin()) {
            admin.describeCluster().nodes().get().forEach(node -> actualBrokerToRackMap.put(node.id(), node.rack()));
        }
        return actualBrokerToRackMap;
    }


    private static ReplicaDistributions getReplicaDistribution(Map<Integer, List<Integer>> assignment,
                                                               Map<Integer, String> brokerRackMapping) {
        Map<Integer, List<String>> partitionToRackMap = new HashMap<>();
        Map<Integer, Integer> partitionToCountMap = new HashMap<>();

        for (Map.Entry<Integer, List<Integer>> entry : assignment.entrySet()) {
            int partitionId = entry.getKey();
            List<Integer> replicaList = entry.getValue();

            for (int brokerId : replicaList) {
                partitionToCountMap.put(brokerId, partitionToCountMap.getOrDefault(brokerId, 0) + 1);
                String rack = brokerRackMapping.get(brokerId);
                if (rack == null) {
                    throw new RuntimeException("No mapping found for " + brokerId + " in `brokerRackMapping`");
                }
                partitionToRackMap.computeIfAbsent(partitionId, k -> new ArrayList<>()).add(rack);
            }
        }
        return new ReplicaDistributions(partitionToRackMap, partitionToCountMap);
    }

    private record ReplicaDistributions(Map<Integer, List<String>> partitionToRackMap,
                                        Map<Integer, Integer> partitionToCountMap) {
    }
}
