package org.ds;

import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.DisplayName;
import java.util.*;

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

@DisplayName("Consistent Hashing Test Suite")
public class ConsistentHashTest {
    
    private ConsistentHash<String> consistentHash;
    
    @BeforeEach
    void setUp() {
        consistentHash = new ConsistentHash<>(150);
    }

    // ============================================
    // Basic Functionality Tests
    // ============================================

    @Test
    @DisplayName("Should handle empty ring")
    void testEmptyRing() {
        assertNull(consistentHash.getNode("anyKey"));
        assertEquals(0, consistentHash.getNodeCount());
        assertEquals(0, consistentHash.getRingSize());
    }

    @Test
    @DisplayName("Should add single node")
    void testAddSingleNode() {
        consistentHash.addNode("node1");
        assertEquals(1, consistentHash.getNodeCount());
        assertEquals(150, consistentHash.getRingSize());
        assertTrue(consistentHash.containsNode("node1"));
    }

    @Test
    @DisplayName("Should add multiple nodes")
    void testAddMultipleNodes() {
        consistentHash.addNode("node1");
        consistentHash.addNode("node2");
        consistentHash.addNode("node3");
        
        assertEquals(3, consistentHash.getNodeCount());
        assertEquals(450, consistentHash.getRingSize());
        assertTrue(consistentHash.getNodes().contains("node1"));
        assertTrue(consistentHash.getNodes().contains("node2"));
        assertTrue(consistentHash.getNodes().contains("node3"));
    }

    @Test
    @DisplayName("Should route key to single node")
    void testGetNodeWithSingleNode() {
        consistentHash.addNode("node1");
        assertEquals("node1", consistentHash.getNode("key1"));
        assertEquals("node1", consistentHash.getNode("key2"));
        assertEquals("node1", consistentHash.getNode("key3"));
    }

    @Test
    @DisplayName("Should route keys consistently")
    void testKeyConsistency() {
        consistentHash.addNode("node1");
        consistentHash.addNode("node2");
        
        String node1 = consistentHash.getNode("testKey");
        String node2 = consistentHash.getNode("testKey");
        String node3 = consistentHash.getNode("testKey");
        
        assertEquals(node1, node2);
        assertEquals(node2, node3);
    }

    // ============================================
    // Node Removal Tests
    // ============================================

    @Test
    @DisplayName("Should remove node from ring")
    void testRemoveNode() {
        consistentHash.addNode("node1");
        consistentHash.addNode("node2");
        
        assertEquals(2, consistentHash.getNodeCount());
        
        consistentHash.removeNode("node1");
        
        assertEquals(1, consistentHash.getNodeCount());
        assertEquals(150, consistentHash.getRingSize());
        assertFalse(consistentHash.containsNode("node1"));
        assertTrue(consistentHash.containsNode("node2"));
    }

    @Test
    @DisplayName("Should route to remaining node after removal")
    void testRoutingAfterRemoval() {
        consistentHash.addNode("node1");
        consistentHash.addNode("node2");
        
        consistentHash.removeNode("node1");
        
        assertEquals("node2", consistentHash.getNode("anyKey"));
    }

    @Test
    @DisplayName("Should remove all nodes")
    void testRemoveAllNodes() {
        consistentHash.addNode("node1");
        consistentHash.addNode("node2");
        consistentHash.addNode("node3");
        
        consistentHash.removeNode("node1");
        consistentHash.removeNode("node2");
        consistentHash.removeNode("node3");
        
        assertEquals(0, consistentHash.getNodeCount());
        assertNull(consistentHash.getNode("anyKey"));
    }

    // ============================================
    // Load Distribution Tests
    // ============================================

    @Test
    @DisplayName("Should distribute keys somewhat evenly with multiple nodes")
    void testLoadDistribution() {
        consistentHash.addNode("node1");
        consistentHash.addNode("node2");
        consistentHash.addNode("node3");
        
        Map<String, Integer> distribution = consistentHash.getDistribution(3000);
        
        // Each node should get some keys
        assertTrue(distribution.get("node1") > 0);
        assertTrue(distribution.get("node2") > 0);
        assertTrue(distribution.get("node3") > 0);
        
        // Total should equal sample size
        int total = distribution.values().stream().mapToInt(Integer::intValue).sum();
        assertEquals(3000, total);
        
        // With 150 virtual nodes, distribution should be relatively balanced
        // Allow 40% deviation from perfect distribution
        int idealCount = 3000 / 3;
        distribution.values().forEach(count -> {
            assertTrue(count > idealCount * 0.6, "Count " + count + " is too low");
            assertTrue(count < idealCount * 1.4, "Count " + count + " is too high");
        });
    }

    @Test
    @DisplayName("Should improve distribution with more virtual nodes")
    void testVirtualNodesImpactDistribution() {
        ConsistentHash<String> hashWith10Nodes = new ConsistentHash<>(10);
        ConsistentHash<String> hashWith500Nodes = new ConsistentHash<>(500);
        
        // Add same nodes to both
        for (ConsistentHash<String> hash : Arrays.asList(hashWith10Nodes, hashWith500Nodes)) {
            hash.addNode("node1");
            hash.addNode("node2");
            hash.addNode("node3");
        }
        
        Map<String, Integer> dist10 = hashWith10Nodes.getDistribution(3000);
        Map<String, Integer> dist500 = hashWith500Nodes.getDistribution(3000);
        
        // Calculate standard deviation for each distribution
        double stdDev10 = calculateStdDev(dist10.values());
        double stdDev500 = calculateStdDev(dist500.values());
        
        // More virtual nodes should result in more uniform distribution
        assertTrue(stdDev500 <= stdDev10, 
            "More virtual nodes should improve distribution uniformity");
    }

    @Test
    @DisplayName("Should demonstrate ring size vs node count relationship")
    void testRingSizeVsNodeCount() {
        int virtualNodes = 150;
        ConsistentHash<String> hash = new ConsistentHash<>(virtualNodes);
        
        for (int i = 1; i <= 5; i++) {
            hash.addNode("node" + i);
            assertEquals(i * virtualNodes, hash.getRingSize());
        }
    }

    // ============================================
    // Key Redistribution Tests
    // ============================================

    @Test
    @DisplayName("Should minimize key redistribution on node addition")
    void testKeyRedistributionOnNodeAddition() {
        consistentHash.addNode("node1");
        consistentHash.addNode("node2");
        
        int sampleSize = 10000;
        Map<String, String> beforeAddition = getKeyAssignments(sampleSize);
        
        consistentHash.addNode("node3");
        
        Map<String, String> afterAddition = getKeyAssignments(sampleSize);
        
        // Count keys that moved to different nodes
        int movedKeys = 0;
        for (String key : beforeAddition.keySet()) {
            if (!beforeAddition.get(key).equals(afterAddition.get(key))) {
                movedKeys++;
            }
        }
        
        // With consistent hashing, only ~1/3 of keys should move
        double movedPercentage = (movedKeys * 100.0) / sampleSize;
        assertTrue(movedPercentage < 50, 
            "Too many keys redistributed: " + movedPercentage + "%");
        assertTrue(movedPercentage > 20, 
            "Expected some key redistribution, got: " + movedPercentage + "%");
    }

    @Test
    @DisplayName("Should minimize key redistribution on node removal")
    void testKeyRedistributionOnNodeRemoval() {
        consistentHash.addNode("node1");
        consistentHash.addNode("node2");
        consistentHash.addNode("node3");
        
        int sampleSize = 10000;
        Map<String, String> beforeRemoval = getKeyAssignments(sampleSize);
        
        consistentHash.removeNode("node3");
        
        Map<String, String> afterRemoval = getKeyAssignments(sampleSize);
        
        // Count keys that moved to different nodes
        int movedKeys = 0;
        for (String key : beforeRemoval.keySet()) {
            if (!beforeRemoval.get(key).equals(afterRemoval.get(key))) {
                movedKeys++;
            }
        }
        
        // Keys that weren't on node3 shouldn't move
        // Keys on node3 should be redistributed among node1 and node2
        double movedPercentage = (movedKeys * 100.0) / sampleSize;
        assertTrue(movedPercentage < 50, 
            "Too many keys redistributed: " + movedPercentage + "%");
    }

    @Test
    @DisplayName("Should handle rapid node additions")
    void testRapidNodeAdditions() {
        for (int i = 1; i <= 10; i++) {
            consistentHash.addNode("node" + i);
        }
        
        assertEquals(10, consistentHash.getNodeCount());
        assertEquals(1500, consistentHash.getRingSize());
        
        // All keys should still be routable
        for (int i = 0; i < 100; i++) {
            assertNotNull(consistentHash.getNode("key:" + i));
        }
    }

    @Test
    @DisplayName("Should handle rapid node removals")
    void testRapidNodeRemovals() {
        for (int i = 1; i <= 10; i++) {
            consistentHash.addNode("node" + i);
        }
        
        for (int i = 10; i > 0; i--) {
            consistentHash.removeNode("node" + i);
        }
        
        assertEquals(0, consistentHash.getNodeCount());
        assertNull(consistentHash.getNode("anyKey"));
    }

    // ============================================
    // Edge Cases and Boundary Tests
    // ============================================

    @Test
    @DisplayName("Should handle duplicate node additions")
    void testDuplicateNodeAddition() {
        consistentHash.addNode("node1");
        int sizeAfterFirst = consistentHash.getRingSize();
        
        // Adding same node again should overwrite (idempotent)
        consistentHash.addNode("node1");
        int sizeAfterSecond = consistentHash.getRingSize();
        
        assertEquals(sizeAfterFirst, sizeAfterSecond);
    }

    @Test
    @DisplayName("Should handle removing non-existent node")
    void testRemoveNonExistentNode() {
        consistentHash.addNode("node1");
        int sizeBeforeRemoval = consistentHash.getRingSize();
        
        // Should not throw exception
        consistentHash.removeNode("node2");
        
        assertEquals(sizeBeforeRemoval, consistentHash.getRingSize());
        assertEquals(1, consistentHash.getNodeCount());
    }

    @Test
    @DisplayName("Should handle empty string keys")
    void testEmptyStringKey() {
        consistentHash.addNode("node1");
        assertNotNull(consistentHash.getNode(""));
    }

    @Test
    @DisplayName("Should handle very long keys")
    void testVeryLongKey() {
        consistentHash.addNode("node1");
        String longKey = "k".repeat(10000);
        assertNotNull(consistentHash.getNode(longKey));
    }

    @Test
    @DisplayName("Should handle special characters in keys")
    void testSpecialCharacterKeys() {
        consistentHash.addNode("node1");
        
        String[] specialKeys = {
            "key-with-dashes",
            "key_with_underscores",
            "key.with.dots",
            "key@with#special$chars",
            "键with中文",
            "🔑emoji🔑"
        };
        
        for (String key : specialKeys) {
            assertNotNull(consistentHash.getNode(key), "Failed for key: " + key);
        }
    }

    @Test
    @DisplayName("Should handle special characters in node names")
    void testSpecialCharacterNodeNames() {
        String[] nodeNames = {
            "server-1",
            "server_2",
            "192.168.1.1",
            "服务器1",
            "🖥️server"
        };
        
        for (String nodeName : nodeNames) {
            assertDoesNotThrow(() -> consistentHash.addNode(nodeName));
        }
        
        assertEquals(nodeNames.length, consistentHash.getNodeCount());
    }

    // ============================================
    // Virtual Node Tests
    // ============================================

    @Test
    @DisplayName("Should respect virtual node count")
    void testVirtualNodeCount() {
        int[] virtualNodeCounts = {1, 10, 50, 150, 500};
        
        for (int count : virtualNodeCounts) {
            ConsistentHash<String> hash = new ConsistentHash<>(count);
            hash.addNode("node1");
            assertEquals(count, hash.getRingSize());
        }
    }

    @Test
    @DisplayName("Should use virtual nodes for distribution")
    void testVirtualNodesDistribution() {
        ConsistentHash<String> hash1 = new ConsistentHash<>(1);
        ConsistentHash<String> hash150 = new ConsistentHash<>(150);
        
        hash1.addNode("node1");
        hash1.addNode("node2");
        
        hash150.addNode("node1");
        hash150.addNode("node2");
        
        Map<String, Integer> dist1 = hash1.getDistribution(1000);
        Map<String, Integer> dist150 = hash150.getDistribution(1000);
        
        // More virtual nodes should be closer to 50-50 split
        int diff1 = Math.abs(dist1.get("node1") - dist1.get("node2"));
        int diff150 = Math.abs(dist150.get("node1") - dist150.get("node2"));
        
        assertTrue(diff150 < diff1, "More virtual nodes should improve distribution balance");
    }

    // ============================================
    // Hash Function Tests
    // ============================================

    @Test
    @DisplayName("Should work with simple hash function")
    void testSimpleHashFunction() {
        ConsistentHash<String> hash = new ConsistentHash<>(150, 
            new ConsistentHash.SimpleHashFunction());
        
        hash.addNode("node1");
        hash.addNode("node2");
        
        assertNotNull(hash.getNode("key1"));
        assertEquals(2, hash.getNodeCount());
    }

    @Test
    @DisplayName("Should work with custom hash function")
    void testCustomHashFunction() {
        ConsistentHash<String> hash = new ConsistentHash<>(150, 
            key -> Math.abs((long) key.hashCode() * 31));
        
        hash.addNode("node1");
        
        assertNotNull(hash.getNode("key1"));
    }

    @Test
    @DisplayName("MD5 hash should be consistent")
    void testMD5HashConsistency() {
        ConsistentHash.HashFunction md5 = new ConsistentHash.MD5HashFunction();
        
        long hash1 = md5.hash("testKey");
        long hash2 = md5.hash("testKey");
        
        assertEquals(hash1, hash2);
    }

    @Test
    @DisplayName("Different keys should have different hashes")
    void testHashDiversity() {
        ConsistentHash.HashFunction md5 = new ConsistentHash.MD5HashFunction();
        
        long hash1 = md5.hash("key1");
        long hash2 = md5.hash("key2");
        long hash3 = md5.hash("key3");
        
        assertNotEquals(hash1, hash2);
        assertNotEquals(hash2, hash3);
        assertNotEquals(hash1, hash3);
    }

    // ============================================
    // Complex Scenarios
    // ============================================

    @Test
    @DisplayName("Should handle complete lifecycle")
    void testCompleteLifecycle() {
        // Start empty
        assertEquals(0, consistentHash.getNodeCount());
        
        // Add nodes
        for (int i = 1; i <= 5; i++) {
            consistentHash.addNode("server" + i);
        }
        assertEquals(5, consistentHash.getNodeCount());
        
        // Keys should be distributed
        Map<String, Integer> dist = consistentHash.getDistribution(1000);
        assertTrue(dist.values().stream().allMatch(v -> v > 0));
        
        // Remove some nodes
        consistentHash.removeNode("server1");
        consistentHash.removeNode("server3");
        assertEquals(3, consistentHash.getNodeCount());
        
        // Keys still routable
        for (int i = 0; i < 100; i++) {
            assertNotNull(consistentHash.getNode("key:" + i));
        }
        
        // Clear all
        consistentHash.clear();
        assertEquals(0, consistentHash.getNodeCount());
        assertNull(consistentHash.getNode("anyKey"));
    }

    @Test
    @DisplayName("Should simulate cache cluster scenario")
    void testCacheClusterScenario() {
        // Initial cluster
        consistentHash.addNode("cache1");
        consistentHash.addNode("cache2");
        consistentHash.addNode("cache3");
        
        // Assign some data
        Map<String, String> assignments = new HashMap<>();
        for (int i = 0; i < 100; i++) {
            String key = "data:" + i;
            assignments.put(key, consistentHash.getNode(key));
        }
        
        // Server failure - remove cache2
        consistentHash.removeNode("cache2");
        
        // Check what was reassigned
        int reassignedCount = 0;
        for (Map.Entry<String, String> entry : assignments.entrySet()) {
            String newNode = consistentHash.getNode(entry.getKey());
            if (!newNode.equals(entry.getValue())) {
                reassignedCount++;
            }
        }
        
        // Only data that was on cache2 should be reassigned
        assertTrue(reassignedCount > 0, "Some data should be reassigned");
        assertTrue(reassignedCount < 100, "Not all data should be reassigned");
        
        // Recovery - add cache2 back
        consistentHash.addNode("cache2");
        
        Map<String, Integer> finalDistribution = consistentHash.getDistribution(1000);
        assertTrue(finalDistribution.get("cache2") > 0, "Cache2 should serve some requests");
    }

    @Test
    @DisplayName("Should handle load balancing scenario")
    void testLoadBalancingScenario() {
        // Web servers in load balancer
        consistentHash.addNode("web-server-1");
        consistentHash.addNode("web-server-2");
        consistentHash.addNode("web-server-3");
        
        // Distribute requests
        Map<String, Integer> load = consistentHash.getDistribution(10000);
        
        // Each server should handle requests
        load.forEach((server, count) -> {
            assertTrue(count > 0, "Server " + server + " should handle requests");
            double percentage = (count * 100.0) / 10000;
            // Allow some variation (20-50% per server for 3 servers)
            assertTrue(percentage > 20, server + " has too few requests: " + percentage + "%");
        });
    }

    // ============================================
    // Helper Methods
    // ============================================

    private Map<String, String> getKeyAssignments(int count) {
        Map<String, String> assignments = new HashMap<>();
        for (int i = 0; i < count; i++) {
            String key = "key:" + i;
            assignments.put(key, consistentHash.getNode(key));
        }
        return assignments;
    }

    private double calculateStdDev(Collection<Integer> values) {
        if (values.isEmpty()) return 0;
        
        double mean = values.stream().mapToDouble(Integer::doubleValue).average().orElse(0);
        double variance = values.stream()
            .mapToDouble(v -> Math.pow(v - mean, 2))
            .average()
            .orElse(0);
        
        return Math.sqrt(variance);
    }
}

