package org.ds;

import java.util.*;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.nio.charset.StandardCharsets;

/**
 * Consistent Hashing implementation with virtual nodes for better load distribution.
 * 
 * Key features:
 * - MD5-based hash function for uniform distribution
 * - Virtual nodes to improve load balance
 * - Support for dynamic node addition/removal
 * - Minimal key redistribution on topology changes
 */
public class ConsistentHash<T> {
    private static final int DEFAULT_VIRTUAL_NODES = 150;
    protected final SortedMap<Long, T> ring = new TreeMap<>();
    protected final int virtualNodes;
    protected final HashFunction hashFunction;

    /**
     * Create a ConsistentHash with default number of virtual nodes
     */
    public ConsistentHash() {
        this(DEFAULT_VIRTUAL_NODES);
    }

    /**
     * Create a ConsistentHash with specified number of virtual nodes
     * 
     * @param virtualNodes number of virtual nodes per physical node
     */
    public ConsistentHash(int virtualNodes) {
        this(virtualNodes, new MD5HashFunction());
    }

    /**
     * Create a ConsistentHash with custom hash function
     * 
     * @param virtualNodes number of virtual nodes per physical node
     * @param hashFunction custom hash function implementation
     */
    public ConsistentHash(int virtualNodes, HashFunction hashFunction) {
        this.virtualNodes = virtualNodes;
        this.hashFunction = hashFunction;
    }

    /**
     * Add a node to the ring
     * 
     * @param node the node to add
     */
    public void addNode(T node) {
        for (int i = 0; i < virtualNodes; i++) {
            long hash = hashFunction.hash(node.toString() + ":" + i);
            ring.put(hash, node);
        }
    }

    /**
     * Remove a node from the ring
     * 
     * @param node the node to remove
     */
    public void removeNode(T node) {
        for (int i = 0; i < virtualNodes; i++) {
            long hash = hashFunction.hash(node.toString() + ":" + i);
            ring.remove(hash);
        }
    }

    /**
     * Get the node responsible for the given key
     * 
     * @param key the key to look up
     * @return the responsible node, or null if ring is empty
     */
    public T getNode(String key) {
        if (ring.isEmpty()) {
            return null;
        }
        
        long hash = hashFunction.hash(key);
        
        // Find the first node hash that is greater than or equal to key hash
        SortedMap<Long, T> tailMap = ring.tailMap(hash);
        if (!tailMap.isEmpty()) {
            return tailMap.get(tailMap.firstKey());
        }
        

        //ask a qesution here
        // Wrap around: return the first node in the ring
        return ring.get(ring.firstKey());
    }

    /**
     * Get all nodes in the ring
     * 
     * @return set of unique nodes
     */
    public Set<T> getNodes() {
        return new HashSet<>(ring.values());
    }

    /**
     * Get the number of unique physical nodes
     * 
     * @return number of unique nodes
     */
    public int getNodeCount() {
        return getNodes().size();
    }

    /**
     * Check if the ring contains a node
     * 
     * @param node the node to check
     * @return true if node exists in ring
     */
    public boolean containsNode(T node) {
        return getNodes().contains(node);
    }

    /**
     * Get the ring size (total virtual nodes)
     * 
     * @return number of virtual nodes in ring
     */
    public int getRingSize() {
        return ring.size();
    }

    /**
     * Get statistics about key distribution across nodes
     * 
     * @param sampleSize number of keys to sample for distribution analysis
     * @return map of node to count of keys assigned to it
     */
    public Map<T, Integer> getDistribution(int sampleSize) {
        Map<T, Integer> distribution = new HashMap<>();
        for (T node : getNodes()) {
            distribution.put(node, 0);
        }
        
        for (int i = 0; i < sampleSize; i++) {
            T node = getNode("key:" + i);
            distribution.put(node, distribution.get(node) + 1);
        }
        
        return distribution;
    }

    /**
     * Get keys assigned to a specific node
     * 
     * @param node the target node
     * @param sampleSize number of sample keys to check
     * @return list of keys assigned to this node
     */
    public List<String> getKeysForNode(T node, int sampleSize) {
        List<String> keys = new ArrayList<>();
        for (int i = 0; i < sampleSize; i++) {
            String key = "key:" + i;
            if (getNode(key).equals(node)) {
                keys.add(key);
            }
        }
        return keys;
    }

    /**
     * Get the node that precedes the given node in the ring
     * 
     * @param node the target node
     * @return the preceding node
     */
    public T getPrecedingNode(T node) {
        Set<T> nodes = getNodes();
        if (nodes.size() <= 1) {
            return null;
        }
        
        // Find the hash positions of all instances of this node
        long lastHash = Long.MIN_VALUE;
        for (Map.Entry<Long, T> entry : ring.entrySet()) {
            if (entry.getValue().equals(node)) {
                lastHash = entry.getKey();
            }
        }
        
        // Find the preceding hash position
        SortedMap<Long, T> headMap = ring.headMap(lastHash);
        if (!headMap.isEmpty()) {
            return headMap.get(headMap.lastKey());
        }
        
        return ring.get(ring.lastKey());
    }

    /**
     * Clear all nodes from the ring
     */
    public void clear() {
        ring.clear();
    }

    /**
     * Interface for hash functions
     */
    public interface HashFunction {
        long hash(String key);
    }

    /**
     * MD5-based hash function for consistent hashing
     */
    public static class MD5HashFunction implements HashFunction {
        @Override
        public long hash(String key) {
            try {
                MessageDigest md = MessageDigest.getInstance("MD5");
                byte[] messageDigest = md.digest(key.getBytes(StandardCharsets.UTF_8));
                
                // Convert first 8 bytes of MD5 to long (unsigned)
                long hash = 0;
                for (int i = 0; i < 8; i++) {
                    hash = (hash << 8) | (messageDigest[i] & 0xFF);
                }
                
                return Math.abs(hash);
            } catch (NoSuchAlgorithmException e) {
                throw new RuntimeException("MD5 algorithm not available", e);
            }
        }
    }

    /**
     * Simple hash function based on String.hashCode() - less uniform but faster
     */
    public static class SimpleHashFunction implements HashFunction {
        @Override
        public long hash(String key) {
            return Math.abs((long) key.hashCode());
        }
    }

    /**
     * Get detailed statistics about the hash ring
     * 
     * @return formatted statistics string
     */
    public String getStatistics(int sampleSize) {
        StringBuilder sb = new StringBuilder();
        sb.append("Consistent Hash Statistics:\n");
        sb.append("============================\n");
        sb.append("Physical Nodes: ").append(getNodeCount()).append("\n");
        sb.append("Virtual Nodes: ").append(getRingSize()).append("\n");
        sb.append("Virtual Nodes per Physical Node: ").append(virtualNodes).append("\n");
        sb.append("\nKey Distribution (based on ").append(sampleSize).append(" samples):\n");
        
        Map<T, Integer> distribution = getDistribution(sampleSize);
        distribution.entrySet().stream()
            .sorted((a, b) -> b.getValue().compareTo(a.getValue()))
            .forEach(entry -> {
                double percentage = (entry.getValue() * 100.0) / sampleSize;
                sb.append("  ").append(entry.getKey()).append(": ")
                  .append(entry.getValue()).append(" (").append(String.format("%.2f%%", percentage)).append(")\n");
            });
        
        return sb.toString();
    }
}

