package com.muyu.mqtt;

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

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

public class CustomWeightedRoundRobinPartitioner implements Partitioner {

    private int currentPartitionIndex = 0;
    private final int[] partitionWeights;

    public CustomWeightedRoundRobinPartitioner(int[] partitionWeights) {
        this.partitionWeights = partitionWeights;
    }

    @Override
    public int partition(String topic, Object key, byte[] keyBytes, Object value, byte[] valueBytes, Cluster cluster) {
        List<PartitionInfo> partitions = cluster.availablePartitionsForTopic(topic);
        int numPartitions = partitions.size();
        if (numPartitions == 0) {
            throw new IllegalStateException("No partitions available for topic " + topic);
        }
        
        // 确保权重数组长度与分区数匹配
        if (partitionWeights.length != numPartitions) {
            throw new IllegalArgumentException("Partition weights length does not match number of partitions");
        }
        
        synchronized (CustomWeightedRoundRobinPartitioner.class) {
            int totalWeight = 0;
            for (int weight : partitionWeights) {
                totalWeight += weight;
            }
            currentPartitionIndex = (currentPartitionIndex + 1) % totalWeight; // 轮询总权重
            
            int chosenPartition = 0;
            for (int i = 0; i < numPartitions; i++) {
                if (currentPartitionIndex < partitionWeights[i]) {
                    chosenPartition = i;
                    break;
                } else {
                    currentPartitionIndex -= partitionWeights[i];
                }
            }
            return chosenPartition;
        }
    }

    @Override
    public void close() {}

    @Override
    public void configure(Map<String, ?> configs) {}
}
