import random
import unittest

from jump_consistent_hash import jump_hash, JumpConsistentHash, _to_u64


class TestJumpHash(unittest.TestCase):
    def test_golden_vectors(self):
        # Golden vectors, produced by jump_ref.c in this directory: the
        # Lamping & Veach figure-5 reference, compiled with `gcc -O2`, fed the
        # key and bucket count and asked for the bucket.  Regenerate with
        #     gcc -O2 -o /tmp/jump_ref jump_ref.c && echo "0 5" | /tmp/jump_ref
        # These pin the MAPPING, not its statistics: a hash can be uniform,
        # minimal-movement and deterministic while sending every key to a
        # different bucket than the reference does.  The last two cells pin a
        # string key and its md5->u64 fold, because the fold is this module's
        # own convention (the paper defines the algorithm on a u64) and it is
        # the line that can move ~99% of keys without any property cell
        # noticing.
        self.assertEqual(jump_hash(0, 5), 0)
        self.assertEqual(jump_hash(18446744073709551615, 100), 92)
        self.assertEqual(_to_u64("key-0"), 18183421273640453746)
        self.assertEqual(jump_hash("key-0", 100), 40)
        self.assertEqual(jump_hash("order-42", 1000), 667)

    def test_range(self):
        for n in (1, 2, 5, 10, 100, 1000):
            for i in range(200):
                b = jump_hash(f"k{i}", n)
                self.assertGreaterEqual(b, 0)
                self.assertLess(b, n)

    def test_single_bucket(self):
        for i in range(50):
            self.assertEqual(jump_hash(i, 1), 0)

    def test_deterministic(self):
        self.assertEqual(jump_hash("abc", 7), jump_hash("abc", 7))

    def test_invalid(self):
        with self.assertRaises(ValueError):
            jump_hash("x", 0)
        with self.assertRaises(ValueError):
            jump_hash("x", -1)
        with self.assertRaises(TypeError):
            jump_hash("x", 1.5)

    def test_minimal_movement(self):
        # Moving from k to k+1 buckets: only ~1/(k+1) keys should move.
        keys = [f"key-{i}" for i in range(10000)]
        for k in (4, 8, 16, 32):
            before = {key: jump_hash(key, k) for key in keys}
            after = {key: jump_hash(key, k + 1) for key in keys}
            moved = sum(1 for key in keys if before[key] != after[key])
            # expected ~ 10000/(k+1); allow generous slack
            expected = 10000 / (k + 1)
            self.assertLess(moved, expected * 2.5, f"k={k} moved={moved}")
            # ... and a lower bound, so a sticky hash (one that never moves a
            # key, and therefore satisfies every "nothing extra moved" test)
            # is graded here rather than only by uniformity.  Measured ratios
            # on this key set are 0.974-1.047, so 0.5 is slack, not a knob.
            self.assertGreater(moved, expected * 0.5,
                               f"k={k} moved={moved} -- below half the "
                               f"expected {expected:.0f}: keys are not moving")

    def test_uniformity(self):
        n = 10
        counts = [0] * n
        for i in range(100000):
            counts[jump_hash(i, n)] += 1
        # each bucket within 5% of mean
        mean = 100000 / n
        for c in counts:
            self.assertGreater(c, mean * 0.9)
            self.assertLess(c, mean * 1.1)

    def test_uniformity_other_bucket_count(self):
        # A second bucket count, so a hash tuned to n=10 (or a bug that only
        # bites when n is not a round number) is not silently exempt.  n=97 is
        # prime and non-power-of-two.  Band: +/-10% of the mean, measured
        # spread on this key set is -5.28%/+4.57% (deterministic, fixed keys),
        # so the band has ~2x headroom and still fails any gross bias.
        n = 97
        keys = 200000
        counts = [0] * n
        for i in range(keys):
            counts[jump_hash(i, n)] += 1
        mean = keys / n
        self.assertEqual(sum(counts), keys)
        for b, c in enumerate(counts):
            self.assertGreater(c, mean * 0.90, f"bucket {b} under-filled: {c}")
            self.assertLess(c, mean * 1.10, f"bucket {b} over-filled: {c}")

    def test_wrapper(self):
        jch = JumpConsistentHash(5)
        self.assertEqual(jch.bucket("a"), jump_hash("a", 5))
        jch.resize(7)
        self.assertEqual(jch.num_buckets, 7)


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