package org.ds;

import java.util.*;

/**
 * Extension: Implements "Heterogeneous Node Weighting".
 * Allows servers with more CPU/RAM to handle a larger share of the ring.
 */
public class WeightedConsistentHash<T> extends ConsistentHash<T> {

    // Store weights to ensure consistent removal
    private final Map<T, Integer> nodeWeights = new HashMap<>();

    public WeightedConsistentHash(int baseVirtualNodes) {
        super(baseVirtualNodes);
    }

    /**
     * Add a node with a specific capacity weight.
     * @param node The physical node
     * @param weight Capacity multiplier (e.g., 1 for small, 5 for large)
     */
    public void addNode(T node, int weight) {
        if (weight <= 0) weight = 1;
        nodeWeights.put(node, weight);
        
        // The "Secret Sauce": Total Virtual Nodes = Base VN * Weight
        int totalVirtualNodesForThisNode = virtualNodes * weight;
        
        for (int i = 0; i < totalVirtualNodesForThisNode; i++) {
            long hash = hashFunction.hash(node.toString() + ":" + i);
            ring.put(hash, node);
        }
    }

    /**
     * Standard addNode defaults to weight 1
     */
    @Override
    public void addNode(T node) {
        addNode(node, 1);
    }

    @Override
    public void removeNode(T node) {
        Integer weight = nodeWeights.getOrDefault(node, 1);
        int totalVirtualNodes = virtualNodes * weight;
        
        for (int i = 0; i < totalVirtualNodes; i++) {
            long hash = hashFunction.hash(node.toString() + ":" + i);
            ring.remove(hash);
        }
        nodeWeights.remove(node);
    }
}