// A hash ring maps resource names to nodes.  Each node owns a set of control
// points (weighted tokens) spread around a ring of `range` positions; the node
// owning the first control point at or after `hash(name) % range` handles the
// resource.  Adding or removing a node only moves the resources it owns, not
// the whole ring, which is what makes the mapping stable.
//
// Insertion is O(n log n) for any number of nodes, resource names are hashed
// with MurmurHash3 over their UTF-16 code units, and the lookup path walks a
// sorted array of control points with an approximate binary search.

///|
/// How control points are placed when `add` is not given an explicit `points`
/// array.
pub(all) enum Distribution {
  /// Scatter points at random positions, retrying on collisions.
  Random
  /// Space points evenly, computed from the number of nodes that need points.
  ///
  /// Uniform placement is deterministic: two rings that add the same nodes in
  /// the same order get identical control points.  Every node receives the
  /// ring's default weight, so per-node `weight` arguments are ignored.
  Uniform
}

///|
/// A consistent hashing ring of nodes of type `T`.
///
/// `T` must implement `Eq` so that a node can be found again by `remove` and
/// de-duplicated by `get_many`; `String` is the usual choice.
///
/// Nodes are stored in `nodes` in insertion order; each node owns one array of
/// control points in `node_keys`.  `key_map` resolves a control point to its
/// node.  The sorted arrays `keys` and `nodes_sorted` are the lookup index and
/// are lazily (re)built by `add`, `remove` and `get`.
///
/// ```mbt check
/// test {
///   let hr : @hashwheel.ConsistentHash[String] = @hashwheel.ConsistentHash::new()
///   ignore(hr.add("server1").add("server2"))
///   let node = hr.get("resourceName")
///   assert_true(node is Some("server1" | "server2"))
/// }
/// ```
pub struct ConsistentHash[T] {
  /// nodes placed on the ring, in the order they were added
  priv nodes : Array[T]
  /// the control points of each node, parallel to `nodes`
  priv node_keys : Array[Array[Int]]
  /// control point -> owning node
  priv mut key_map : Map[Int, T]
  /// every control point currently reserved or in use
  priv mut point_set : Map[Int, Unit]
  /// sorted control points, or `None` when they need rebuilding
  priv mut keys : Array[Int]?
  /// node owning each entry of `keys`, parallel to it
  priv mut nodes_sorted : Array[T]?
  priv cache : LruCache[T]?
  priv mut need_key_map : Bool
  priv range : Int
  priv weight : Int
  priv uniform : Bool
  priv order_nodes : ((Array[T]) -> Array[T])?
  priv rng : @random.Rand
  /// number of nodes on the ring
  mut node_count : Int
  /// number of control points around the ring
  mut key_count : Int
}

///|
/// Default hash ring capacity: the control point modulo.
pub let default_range : Int = 100003

///|
/// Default number of control points created per node.
pub let default_weight : Int = 40

///|
/// Create an empty ring.
///
/// Parameters (all optional, pass by label):
///
/// * `range` : hash ring capacity, the modulus of the sorted control points.
///   Smaller values distribute better; the default is 100003.  Prefer an odd
///   range relatively prime to the number of nodes.
/// * `weight` : default number of control points per node, default 40.  Three
///   nodes added with weights 1, 2 and 3 handle 1/6, 1/3 and 1/2 of the
///   resources each.
/// * `distribution` : `Random` (default) or `Uniform`.
/// * `order_nodes` : called with the nodes that still need uniformly
///   distributed control points; returns them in the order to assign points in.
///   Only used by `Uniform`.
/// * `cache` : size of an optional LRU cache of `get` results, 0 disables it.
///   The cache is dropped whenever a node is added or removed.
/// * `nodes` : nodes to `add` right away.
/// * `seed` : a 32-byte seed for the control point generator; the same seed
///   reproduces the same random ring.
///
/// ```mbt check
/// test {
///   let hr : @hashwheel.ConsistentHash[String] = @hashwheel.ConsistentHash::new(nodes=[
///     "a", "b", "c",
///   ])
///   assert_eq(hr.node_count, 3)
///   assert_eq(hr.get_nodes(), ["a", "b", "c"])
/// }
/// ```
pub fn[T] ConsistentHash::new(
  range? : Int,
  weight? : Int,
  distribution? : Distribution,
  order_nodes? : (Array[T]) -> Array[T],
  cache? : Int,
  nodes? : Array[T],
  seed? : Bytes,
) -> ConsistentHash[T] {
  let range = range.unwrap_or(default_range)
  let weight = weight.unwrap_or(default_weight)
  let distribution = distribution.unwrap_or(Random)
  let cache = cache.unwrap_or(0)
  let rng = match seed {
    Some(s) => @random.Rand::chacha8(seed=s)
    None => @random.Rand::new()
  }
  let ring = ConsistentHash::{
    nodes: [],
    node_keys: [],
    key_map: Map([]),
    point_set: Map([]),
    keys: None,
    nodes_sorted: None,
    cache: if cache > 0 {
      Some(LruCache::new(cache))
    } else {
      None
    },
    need_key_map: false,
    range,
    weight,
    uniform: distribution is Uniform,
    order_nodes,
    rng,
    node_count: 0,
    key_count: 0,
  }
  if nodes is Some(ns) {
    for n in ns {
      ignore(ring.add(n))
    }
  }
  ring
}

///|
/// Register `node` as also managing resources, with `weight` control points
/// (or an explicit `points` array).
///
/// The node's share of the resources is proportionate to its weight; the
/// default is the ring's weight.  Adding the same node twice increases its
/// weight.  Returns the ring, so calls can be chained.
///
/// ```mbt check
/// test {
///   let hr : @hashwheel.ConsistentHash[String] = @hashwheel.ConsistentHash::new(
///     range=24,
///     weight=4,
///   )
///   ignore(hr.add("a").add("b", weight=8))
///   assert_eq(hr.key_count, 12)
/// }
/// ```
pub fn[T] ConsistentHash::add(
  self : ConsistentHash[T],
  node : T,
  weight? : Int,
  points? : Array[Int],
) -> ConsistentHash[T] {
  let node_weight = match weight {
    Some(w) => w
    None => self.weight
  }
  let pts = match points {
    Some(p) => p.copy()
    None =>
      if self.uniform {
        // leave placeholders; build_key_map assigns the points on first use
        self.need_key_map = true
        Array::make(node_weight, -1)
      } else {
        self.make_control_points(node_weight)
      }
  }
  self.nodes.push(node)
  self.node_keys.push(pts)
  if pts.length() > 0 && pts[0] >= 0 {
    self.map_node_points(node, pts)
  }
  self.keys = None
  self.nodes_sorted = None
  if self.cache is Some(c) {
    c.clear()
  }
  self.key_count = self.key_count + pts.length()
  self.node_count = self.node_count + 1
  self
}

///|
/// Remove all entries of `node` from the ring and free its control points.
/// Freed points may be handed to nodes added later.  Removing a node that is
/// not on the ring does nothing.  Returns the ring.
///
/// ```mbt check
/// test {
///   let hr : @hashwheel.ConsistentHash[String] = @hashwheel.ConsistentHash::new(nodes=[
///     "a", "b", "a",
///   ])
///   assert_eq(hr.node_count, 3)
///   ignore(hr.remove("a"))
///   assert_eq(hr.node_count, 1)
///   assert_eq(hr.get_nodes(), ["b"])
/// }
/// ```
pub fn[T : Eq] ConsistentHash::remove(
  self : ConsistentHash[T],
  node : T,
) -> ConsistentHash[T] {
  while true {
    match self.nodes.search(node) {
      None => break
      Some(ix) => {
        let pts = self.node_keys[ix]
        let last = self.nodes.length() - 1
        self.nodes[ix] = self.nodes[last]
        ignore(self.nodes.pop())
        self.node_keys[ix] = self.node_keys[last]
        ignore(self.node_keys.pop())
        self.keys = None
        self.nodes_sorted = None
        if self.cache is Some(c) {
          c.clear()
        }
        // the freed points must not stay reserved
        self.need_key_map = true
        self.key_map = Map([])
        self.point_set = Map([])
        self.node_count = self.node_count - 1
        self.key_count = self.key_count - pts.length()
      }
    }
  }
  self
}

///|
/// Return the node that handles `name`, or `None` if the ring is empty.
///
/// ```mbt check
/// test {
///   let hr : @hashwheel.ConsistentHash[String] = @hashwheel.ConsistentHash::new(nodes=[
///     "n1",
///   ])
///   assert_eq(hr.get("anything"), Some("n1"))
/// }
/// ```
pub fn[T] ConsistentHash::get(self : ConsistentHash[T], name : String) -> T? {
  if self.key_count == 0 {
    return None
  }
  if self.cache is Some(c) {
    match c.get(name) {
      Some(node) => return Some(node)
      None => ()
    }
  }
  let index = self.locate(name)
  let node = self.nodes_sorted.unwrap()[index]
  if self.cache is Some(c) {
    c.set(name, node)
  }
  Some(node)
}

///|
/// Return up to `count` distinct nodes, starting with the node that handles
/// `name` and continuing to the nearest following nodes around the ring.
///
/// Returns `None` if the ring is empty.  If fewer distinct nodes than `count`
/// are on the ring, all of them are returned.  Never uses the lookup cache.
///
/// ```mbt check
/// test {
///   let hr : @hashwheel.ConsistentHash[String] = @hashwheel.ConsistentHash::new(nodes=[
///     "a", "b", "c",
///   ])
///   let nodes = hr.get_many("resource", 2)
///   assert_eq(nodes.unwrap().length(), 2)
/// }
/// ```
pub fn[T : Eq] ConsistentHash::get_many(
  self : ConsistentHash[T],
  name : String,
  count : Int,
) -> Array[T]? {
  if self.key_count == 0 {
    return None
  }
  let index = self.locate(name)
  let sorted = self.nodes_sorted.unwrap()
  let found : Array[T] = []
  for i = index; i < self.key_count; i = i + 1 {
    if found.length() >= count {
      break
    }
    let node = sorted[i]
    if !found.contains(node) {
      found.push(node)
    }
  }
  for i = 0; i < index; i = i + 1 {
    if found.length() >= count {
      break
    }
    let node = sorted[i]
    if !found.contains(node) {
      found.push(node)
    }
  }
  Some(found)
}

///|
/// All nodes currently on the ring, in insertion order.
pub fn[T] ConsistentHash::get_nodes(self : ConsistentHash[T]) -> Array[T] {
  self.nodes.copy()
}

///|
/// The control points assigned to `node`, or `None` if the node is not on the
/// ring.
pub fn[T : Eq] ConsistentHash::get_points(
  self : ConsistentHash[T],
  node : T,
) -> Array[Int]? {
  if self.need_key_map {
    self.build_key_map()
  }
  match self.nodes.search(node) {
    None => None
    Some(ix) => Some(self.node_keys[ix].copy())
  }
}

///|
/// Index of the first control point at or after `name` on the ring.
fn[T] ConsistentHash::locate(self : ConsistentHash[T], name : String) -> Int {
  if self.need_key_map {
    self.build_key_map()
  }
  if self.keys is None {
    self.build_keys()
  }
  absearch(
    self.keys.unwrap(),
    (murmur3(name) % self.range.reinterpret_as_uint()).reinterpret_as_int(),
  )
}

///|
/// Create `n` random control points, avoiding the points already on the ring.
///
/// Collision detection is probabilistic (good enough for millions of points);
/// the retry bound applies per point, not per batch, so a batch can fill a ring
/// up to about 90% before giving up.
fn[T] ConsistentHash::make_control_points(
  self : ConsistentHash[T],
  n : Int,
) -> Array[Int] {
  // a prior remove() invalidated the key map; rebuild it before probing for
  // free points
  if self.need_key_map {
    self.build_key_map()
  }
  let points : Array[Int] = []
  for _ in 0..= 100 {
        abort("unable to find an unused control point, tried 100")
      }
    }
    points.push(key)
    // reserve the point so it is not reused, not even by this node
    self.point_set.set(key, ())
  }
  points
}

///|
/// Assign uniformly distributed control points to every node that does not
/// have them yet, and rebuild the control point map.
fn[T] ConsistentHash::build_key_map(self : ConsistentHash[T]) -> Unit {
  // nodes that still need points, by position on the ring
  let new_pos : Array[Int] = []
  for i = 0; i < self.node_keys.length(); i = i + 1 {
    let keys = self.node_keys[i]
    if keys.length() > 0 && keys[0] < 0 {
      new_pos.push(i)
    }
  }
  // optionally reorder the new nodes to make point assignment deterministic
  if self.order_nodes is Some(order) {
    let sorted = order(new_pos.map(i => self.nodes[i]))
    for i = 0; i < new_pos.length(); i = i + 1 {
      self.nodes[new_pos[i]] = sorted[i]
    }
  }
  // spread the points evenly: each new node gets the ring's default weight,
  // and the points of consecutive nodes are interleaved around the ring
  let new_count = new_pos.length()
  let point_count = new_count * self.weight
  let step = self.range.to_double() / point_count.to_double()
  let spacing = step * new_count.to_double()
  for i = 0; i < new_count; i = i + 1 {
    let offset = step * i.to_double() + step / 2.0
    let keys : Array[Int] = []
    for j = 0; j < self.weight; j = j + 1 {
      keys.push((offset + spacing * j.to_double()).round().to_int())
    }
    self.node_keys[new_pos[i]] = keys
  }
  // rebuild the control point -> node map from the live nodes
  self.key_map = Map([])
  self.point_set = Map([])
  let mut total = 0
  for i = 0; i < self.node_keys.length(); i = i + 1 {
    let pts = self.node_keys[i]
    self.map_node_points(self.nodes[i], pts)
    total = total + pts.length()
  }
  self.key_count = total
  self.need_key_map = false
}

///|
/// Regenerate the sorted control point array and the parallel array of owning
/// nodes.
fn[T] ConsistentHash::build_keys(self : ConsistentHash[T]) -> Unit {
  if self.need_key_map {
    self.build_key_map()
  }
  let keys : Array[Int] = []
  for pts in self.node_keys {
    for p in pts {
      keys.push(p)
    }
  }
  keys.sort()
  // parallel array of the node owning each sorted control point, so the hot
  // lookup path can index by position instead of a map lookup
  let nodes : Array[T] = []
  for k in keys {
    nodes.push(self.key_map.at(k))
  }
  self.keys = Some(keys)
  self.nodes_sorted = Some(nodes)
  self.key_count = keys.length()
}

///|
/// Record the node owning each of `points`, and mark the points as used.
fn[T] ConsistentHash::map_node_points(
  self : ConsistentHash[T],
  node : T,
  points : Array[Int],
) -> Unit {
  for p in points {
    self.key_map.set(p, node)
    self.point_set.set(p, ())
  }
}

///|
/// Index of the first value in the sorted `array` that is `>= key`, or 0 if
/// `key` is larger than every value, or -1 for an empty array.
///
/// Binary search until the remaining range is small, then linear search: that
/// is faster than binary searching all the way down for small arrays.
fn absearch(array : Array[Int], key : Int) -> Int {
  let len = array.length()
  let gap = 25
  let mut i = 0
  let mut j = len - 1
  while j - i > gap {
    let mid = (i + j) >> 1
    if array[mid] < key {
      i = mid + 1
    } else {
      j = mid
    }
  }
  while i < len {
    if array[i] >= key {
      return i
    }
    i = i + 1
  }
  if len == 0 {
    -1
  } else {
    0
  }
}