package com.muyu.kafka;

import org.apache.kafka.clients.producer.Partitioner;
import org.apache.kafka.common.Cluster;
import org.apache.kafka.common.PartitionInfo;
import org.springframework.stereotype.Component;

import java.util.List;
import java.util.Map;

/**
 * 自定义分区策略
 */
@Component
public class CustomPartitioner implements Partitioner {
//    Random random = new Random();
    // topic = topic0   key = vin  轮询分区
    private int currentPartitionIndex = 0;
    private final int[] partitionWeights = {1, 1, 1, 1, 1, 1, 1, 1}; // 分区权重

    public int partition(String topic, Object key, byte[] keyBytes, Object value, byte[] valueBytes, Cluster cluster) {
        List<PartitionInfo> partitionInfos = cluster.partitionsForTopic(topic);
        int numPartitions = partitionInfos.size();

        int totalWeight = 0;
        for (int weight : partitionWeights) {
            totalWeight += weight;
        }

        synchronized (this) {
            int chosenPartition = currentPartitionIndex;
            currentPartitionIndex = (currentPartitionIndex + 1) % numPartitions;
            // 根据权重选择分区
            int cumulativeWeight = 0;
            for (int i = 0; i < numPartitions; i++) {
                cumulativeWeight += partitionWeights[i];
                if (chosenPartition < cumulativeWeight) {
                    chosenPartition = i;
                    break;
                }
            }
            System.out.println("当前Key:" + key + "-----> 当前Value:" + value + "----->" + "当前存储分区:" + chosenPartition);
            return chosenPartition;
        }
    }


    public void close() {

    }

    public void configure(Map<String, ?> map) {

    }
}
