"""Jump consistent hash (Lamping & Veach, 2014).

Maps a 64-bit key to a bucket in [0, num_buckets) with minimal
rebalancing when the bucket count changes: only ~1/(k+1) of keys move
when going from k to k+1 buckets.

The reference algorithm uses 64-bit wrapping arithmetic and a double
for the jump computation. This implementation reproduces it exactly
using Python ints masked to 64 bits, so results match the C reference.
"""
from __future__ import annotations

import hashlib
import struct
from typing import Any, Hashable


_MASK64 = (1 << 64) - 1
_MULT = 2862933555777941757
_SCALE = 1 << 31  # 2^31


def _to_u64(key: Any) -> int:
    if isinstance(key, int):
        return key & _MASK64
    if isinstance(key, bytes):
        raw = key
    else:
        raw = str(key).encode("utf-8")
    # md5 -> 128 bits -> fold to 64
    digest = hashlib.md5(raw).digest()
    (hi, lo) = struct.unpack(">QQ", digest)
    return (hi ^ lo) & _MASK64


def jump_hash(key: Any, num_buckets: int) -> int:
    """Return bucket index in [0, num_buckets) for key."""
    if not isinstance(num_buckets, int) or isinstance(num_buckets, bool):
        raise TypeError("num_buckets must be an int")
    if num_buckets < 0:
        raise ValueError("num_buckets must be >= 0")
    if num_buckets == 0:
        raise ValueError("num_buckets must be >= 1")
    if num_buckets == 1:
        return 0

    k = _to_u64(key)
    b = -1
    j = 0
    while j < num_buckets:
        b = j
        k = (k * _MULT + 1) & _MASK64
        # j = (b+1) * 2^31 / ((key >> 33) + 1)
        denom = (k >> 33) + 1
        j = (b + 1) * _SCALE // denom
    return b


class JumpConsistentHash:
    """Convenience wrapper keeping a bucket count."""

    def __init__(self, num_buckets: int):
        if num_buckets < 1:
            raise ValueError("num_buckets must be >= 1")
        self.num_buckets = num_buckets

    def bucket(self, key: Any) -> int:
        return jump_hash(key, self.num_buckets)

    def resize(self, new_num_buckets: int) -> "JumpConsistentHash":
        if new_num_buckets < 1:
            raise ValueError("num_buckets must be >= 1")
        self.num_buckets = new_num_buckets
        return self