import unittest
from consistent_hash import ConsistentHashRing


class TestConsistentHash(unittest.TestCase):
    def test_empty_ring_returns_none(self):
        r = ConsistentHashRing()
        self.assertIsNone(r.get_node("k"))

    def test_single_node_always(self):
        r = ConsistentHashRing()
        r.add_node("A")
        for k in ["x", "y", "z", 1, 2]:
            self.assertEqual(r.get_node(k), "A")

    def test_add_remove(self):
        r = ConsistentHashRing()
        r.add_node("A"); r.add_node("B")
        self.assertEqual(len(r), 2)
        r.remove_node("A")
        self.assertEqual(len(r), 1)
        self.assertNotIn("A", r)
        for k in ["k1", "k2", "k3"]:
            self.assertEqual(r.get_node(k), "B")

    def test_duplicate_add_idempotent(self):
        r = ConsistentHashRing(num_vnodes=5)
        r.add_node("A"); r.add_node("A")
        self.assertEqual(len(r), 1)
        self.assertEqual(len(r._ring), 5)

    def test_remove_nonexistent_noop(self):
        r = ConsistentHashRing()
        r.add_node("A")
        r.remove_node("Z")
        self.assertEqual(len(r), 1)

    def test_determinism(self):
        r1 = ConsistentHashRing(num_vnodes=10)
        r2 = ConsistentHashRing(num_vnodes=10)
        for n in ["n1", "n2", "n3"]:
            r1.add_node(n); r2.add_node(n)
        for k in ["a", "b", "c", "d"]:
            self.assertEqual(r1.get_node(k), r2.get_node(k))

    def test_minimal_rebalancing(self):
        # adding a node should move only a fraction of keys
        r = ConsistentHashRing(num_vnodes=50)
        for n in ["n1", "n2"]:
            r.add_node(n)
        before = {k: r.get_node(k) for k in [f"k{i}" for i in range(200)]}
        r.add_node("n3")
        moved = sum(1 for k in before if r.get_node(k) != before[k])
        # expect roughly 1/3 moved, definitely not all
        self.assertLess(moved, 200)
        self.assertGreater(moved, 0)

    def test_get_nodes_replication(self):
        r = ConsistentHashRing(num_vnodes=20)
        for n in ["a", "b", "c"]:
            r.add_node(n)
        ns = r.get_nodes("key", 2)
        self.assertEqual(len(ns), 2)
        self.assertEqual(len(set(ns)), 2)
        # n larger than node count
        ns2 = r.get_nodes("key", 10)
        self.assertEqual(len(ns2), 3)
        self.assertEqual(len(set(ns2)), 3)

    def test_get_nodes_empty(self):
        r = ConsistentHashRing()
        self.assertEqual(r.get_nodes("k", 3), [])

    def test_invalid_vnodes(self):
        with self.assertRaises(ValueError):
            ConsistentHashRing(0)
        with self.assertRaises(ValueError):
            ConsistentHashRing(-5)

    def test_invalid_n(self):
        r = ConsistentHashRing()
        r.add_node("A")
        with self.assertRaises(ValueError):
            r.get_nodes("k", 0)

    def test_nodes_sorted(self):
        r = ConsistentHashRing()
        for n in ["z", "a", "m"]:
            r.add_node(n)
        self.assertEqual(r.nodes, ["a", "m", "z"])


if __name__ == "__main__":
    unittest.main()