// Copyright 2026 International Digital Economy Academy
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Types

///|
priv struct Entry[K] {
  mut prev : Int
  mut next : Entry[K]?
  mut psl : Int
  hash : Int
  key : K
}

///|
/// Mutable linked hash set that maintains the order of insertion, not thread safe.
///
/// # Example
///
/// ```mbt check
/// test {
///   let set = @set.Set(["three", "eight", "one"])
///   @test.assert_eq(set.contains("two"), false)
///   @test.assert_eq(set.contains("three"), true)
///   set.add("three") // no effect since it already exists
///   set.add("two")
///   @test.assert_eq(set.contains("two"), true)
/// }
/// ```
struct Set[K] {
  mut entries : FixedArray[Entry[K]?]
  mut size : Int // active keys count
  mut capacity : Int // current capacity
  mut capacity_mask : Int // capacity_mask = capacity - 1, used to find idx
  mut grow_at : Int // threshold that triggers grow
  mut head : Entry[K]? // head of linked list
  mut tail : Int // tail of linked list
}

// Implementations

///|
let default_init_capacity = 8

///|
fn[K] new_set(capacity : Int) -> Set[K] {
  let capacity = capacity.next_power_of_two()
  {
    size: 0,
    capacity,
    capacity_mask: capacity - 1,
    grow_at: calc_grow_threshold(capacity),
    entries: FixedArray::make(capacity, None),
    head: None,
    tail: -1,
  }
}

///|
fn capacity_for_length(length : Int) -> Int {
  let mut capacity = length.next_power_of_two()
  if length > calc_grow_threshold(capacity) {
    capacity *= 2
  }
  capacity
}

///|
/// Creates a hash set containing all elements from the given array, preserving insertion order.
/// The optional `capacity` is treated as a minimum initial capacity and will be
/// rounded up to the smallest power of 2 that can hold the requested capacity.
#alias(from_array)
#as_free_fn(from_array)
#alias(of, deprecated="Use from_array instead")
#as_free_fn(of, deprecated="Use from_array instead")
pub fn[K : Hash + Eq] Set::Set(arr : ArrayView[K], capacity? : Int) -> Set[K] {
  let length = arr.length()
  let capacity = match capacity {
    Some(capacity) => capacity.max(capacity_for_length(length))
    None =>
      if length == 0 {
        default_init_capacity
      } else {
        capacity_for_length(length)
      }
  }
  let m = new_set(capacity)
  // arr.each(e => m.add(e))
  // FIXME:(upstream)
  // ArrayView::each depends on array package
  // adding array package caused `unused package` array
  // the root cause is that unused warning is not precise
  for e in arr {
    m.add(e)
  }
  m
}

///|
/// Insert a key into the hash set.
///
/// Parameters:
///
/// * `set` : The hash set to modify.
/// * `key` : The key to insert. Must implement `Hash` and `Eq` traits.
///
/// Example:
///
/// ```mbt check
/// test {
///   let set : @set.Set[String] = Set([])
///   set.add("key")
///   inspect(set.contains("key"), content="true")
///   set.add("key") // no effect since it already exists
///   inspect(set.length(), content="1")
/// }
/// ```
#alias(insert, deprecated)
pub fn[K : Hash + Eq] Set::add(self : Set[K], key : K) -> Unit {
  self.add_with_hash(key, Hash::hash(key))
}

///|
fn[K : Eq] Set::add_with_hash(self : Set[K], key : K, hash : Int) -> Unit {
  if self.size >= self.grow_at {
    self.grow()
  }
  let (idx, psl) = for psl = 0, idx = hash & self.capacity_mask {
    match self.entries[idx] {
      None => break (idx, psl)
      Some(curr_entry) => {
        if curr_entry.hash == hash && curr_entry.key == key {
          return
        }
        if psl > curr_entry.psl {
          self.push_away(idx, curr_entry)
          break (idx, psl)
        }
        continue psl + 1, (idx + 1) & self.capacity_mask
      }
    }
  }
  let entry = { prev: self.tail, next: None, psl, key, hash }
  self.add_entry_to_tail(idx, entry)
}

///|
#owned(entry)
fn[K] Set::push_away(self : Set[K], idx : Int, entry : Entry[K]) -> Unit {
  for psl = entry.psl + 1, idx = (idx + 1) & self.capacity_mask, entry = entry {
    match self.entries[idx] {
      None => {
        entry.psl = psl
        self.set_entry(entry, idx)
        break
      }
      Some(curr_entry) =>
        if psl > curr_entry.psl {
          entry.psl = psl
          self.set_entry(entry, idx)
          continue curr_entry.psl + 1,
            (idx + 1) & self.capacity_mask,
            curr_entry
        } else {
          continue psl + 1, (idx + 1) & self.capacity_mask, entry
        }
    }
  }
}

///|
#owned(entry)
fn[K] Set::set_entry(self : Set[K], entry : Entry[K], new_idx : Int) -> Unit {
  // Fix up the neighbor links before the store: the store consumes the owned
  // `entry` reference, and touching `entry` after it would force the compiler
  // to keep `entry` alive across the store with an extra incref/decref pair.
  match entry.next {
    None => self.tail = new_idx
    Some(next) => next.prev = new_idx
  }
  self.entries[new_idx] = Some(entry)
}

///|
/// Insert a key into the hash set and returns whether the key was successfully added.
///
/// Parameters:
///
/// * `set` : The hash set to modify.
/// * `key` : The key to insert. Must implement `Hash` and `Eq` traits.
///
/// Returns `true` if the key was successfully added (i.e., it wasn't already present),
/// `false` if the key already existed in the set.
///
/// Example:
///
/// ```mbt check
/// test {
///   let set : @set.Set[String] = Set([])
///   inspect(set.add_and_check("key"), content="true") // First insertion
///   inspect(set.add_and_check("key"), content="false") // Already exists
///   inspect(set.length(), content="1")
/// }
/// ```
pub fn[K : Hash + Eq] Set::add_and_check(self : Set[K], key : K) -> Bool {
  if self.size >= self.grow_at {
    self.grow()
  }
  let hash = Hash::hash(key)
  let (idx, psl, added) = for psl = 0, idx = hash & self.capacity_mask {
    match self.entries[idx] {
      None => break (idx, psl, true)
      Some(curr_entry) => {
        if curr_entry.hash == hash && curr_entry.key == key {
          break (idx, psl, false)
        }
        if psl > curr_entry.psl {
          self.push_away(idx, curr_entry)
          break (idx, psl, true)
        }
        continue psl + 1, (idx + 1) & self.capacity_mask
      }
    }
  }
  if added {
    let entry = { prev: self.tail, next: None, psl, key, hash }
    self.add_entry_to_tail(idx, entry)
  }
  added
}

///|
/// Check if the hash set contains a key.
pub fn[K : Hash + Eq] Set::contains(self : Set[K], key : K) -> Bool {
  // inline lookup to avoid unnecessary allocations
  let hash = Hash::hash(key)
  for i = 0, idx = hash & self.capacity_mask {
    guard self.entries[idx] is Some(entry) else { break false }
    if entry.hash == hash && entry.key == key {
      break true
    }
    if i > entry.psl {
      break false
    }
    continue i + 1, (idx + 1) & self.capacity_mask
  }
}

///|
/// Remove a key from the hash set. If the key exists in the set, removes it
/// and adjusts the probe sequence length (PSL) of subsequent entries to
/// maintain the Robin Hood hashing invariant. If the key does not exist,
/// the set remains unchanged.
///
/// Parameters:
///
/// * `self` : The hash set to remove the key from.
/// * `key` : The key to remove from the set.
///
/// Example:
///
/// ```mbt check
/// test {
///   let set = @set.Set(["a", "b"])
///   set.remove("a")
///   inspect(set.contains("a"), content="false")
///   inspect(set.length(), content="1")
/// }
/// ```
pub fn[K : Hash + Eq] Set::remove(self : Set[K], key : K) -> Unit {
  let hash = Hash::hash(key)
  for i = 0, idx = hash & self.capacity_mask {
    guard self.entries[idx] is Some(entry) else { break }
    if entry.hash == hash && entry.key == key {
      self.remove_entry(entry)
      self.shift_back(idx)
      self.size -= 1
      break
    }
    if i > entry.psl {
      break
    }
    continue i + 1, (idx + 1) & self.capacity_mask
  }
}

///|
/// Remove a key from the hash set and returns whether the key was successfully removed.
///
/// Parameters:
///
/// * `set` : The hash set to modify.
/// * `key` : The key to remove. Must implement `Hash` and `Eq` traits.
///
/// Returns `true` if the key was successfully removed (i.e., it was present),
/// `false` if the key didn't exist in the set.
///
/// Example:
///
/// ```mbt check
/// test {
///   let set = @set.Set(["a", "b"])
///   inspect(set.remove_and_check("a"), content="true") // Successfully removed
///   inspect(set.remove_and_check("a"), content="false") // Already removed
///   inspect(set.length(), content="1")
/// }
/// ```
pub fn[K : Hash + Eq] Set::remove_and_check(self : Set[K], key : K) -> Bool {
  let hash = Hash::hash(key)
  for i = 0, idx = hash & self.capacity_mask {
    guard self.entries[idx] is Some(entry) else { break false }
    if entry.hash == hash && entry.key == key {
      self.remove_entry(entry)
      self.shift_back(idx)
      self.size -= 1
      break true
    }
    if i > entry.psl {
      break false
    }
    continue i + 1, (idx + 1) & self.capacity_mask
  }
}

///|
#owned(entry)
fn[K] Set::add_entry_to_tail(
  self : Set[K],
  idx : Int,
  entry : Entry[K],
) -> Unit {
  match self.tail {
    -1 => self.head = Some(entry)
    tail => self.entries[tail].unwrap().next = Some(entry)
  }
  self.tail = idx
  self.entries[idx] = Some(entry)
  self.size += 1
}

///|
fn[K] Set::remove_entry(self : Set[K], entry : Entry[K]) -> Unit {
  match entry.prev {
    -1 => self.head = entry.next
    idx => self.entries[idx].unwrap().next = entry.next
  }
  match entry.next {
    None => self.tail = entry.prev
    Some(next) => next.prev = entry.prev
  }
}

///|
fn[K] Set::shift_back(self : Set[K], idx : Int) -> Unit {
  for cur = idx {
    let next = (cur + 1) & self.capacity_mask
    match self.entries[next] {
      None | Some({ psl: 0, .. }) => {
        self.entries[cur] = None
        break
      }
      Some(entry) => {
        entry.psl -= 1
        self.set_entry(entry, cur)
        continue next
      }
    }
  }
}

///|
fn[K] Set::grow(self : Set[K]) -> Unit {
  let old_head = self.head
  let new_capacity = self.capacity << 1
  self.entries = FixedArray::make(new_capacity, None)
  self.capacity = new_capacity
  self.capacity_mask = new_capacity - 1
  self.grow_at = calc_grow_threshold(self.capacity)
  self.size = 0
  self.head = None
  self.tail = -1
  for x = old_head {
    match x {
      None => break
      Some(e) => {
        let next_in_chain = e.next
        e.next = None
        self.rehash_place_entry(e)
        continue next_in_chain
      }
    }
  }
}

///|
#owned(outer)
fn[K] Set::rehash_place_entry(self : Set[K], outer : Entry[K]) -> Unit {
  let hash = outer.hash
  for psl = 0, idx = hash & self.capacity_mask {
    match self.entries[idx] {
      None => {
        outer.psl = psl
        outer.prev = self.tail
        self.add_entry_to_tail(idx, outer)
        return
      }
      Some(curr) =>
        if psl > curr.psl {
          self.push_away(idx, curr)
          outer.psl = psl
          outer.prev = self.tail
          self.add_entry_to_tail(idx, outer)
          return
        } else {
          continue psl + 1, (idx + 1) & self.capacity_mask
        }
    }
  }
}

// Utils

///|
#deprecated("Use @debug.Debug instead of Show for debugging purposes. See https://github.com/moonbitlang/core/blob/main/debug/README.mbt.md")
pub impl[K : Show] Show for Set[K]

///|
pub impl[K : Show] Show for Set[K] with fn output(self, logger) {
  logger.write_string("{")
  for i = 0, curr = self.head {
    match (i, curr) {
      (_, None) => break logger.write_string("}")
      (i, Some({ key, next, .. })) => {
        if i > 0 {
          logger.write_string(", ")
        }
        logger.write_object(key)
        continue i + 1, next
      }
    }
  }
}

///|
/// Get the number of keys in the set.
#alias(size, deprecated)
pub fn[K] Set::length(self : Set[K]) -> Int {
  self.size
}

///|
/// Get the capacity of the set.
pub fn[K] Set::capacity(self : Set[K]) -> Int {
  self.capacity
}

///|
/// Check if the hash set is empty.
pub fn[K] Set::is_empty(self : Set[K]) -> Bool {
  self.size == 0
}

///|
/// Iterate over all keys of the set in the order of insertion.
#locals(f)
pub fn[K] Set::each(self : Set[K], f : (K) -> Unit raise?) -> Unit raise? {
  for x = self.head {
    match x {
      Some({ key, next, .. }) => {
        f(key)
        continue next
      }
      None => break
    }
  }
}

///|
/// Iterate over all keys of the set in the order of insertion, with index.
#locals(f)
pub fn[K] Set::eachi(self : Set[K], f : (Int, K) -> Unit raise?) -> Unit raise? {
  for i = 0, curr = self.head {
    match (i, curr) {
      (i, Some({ key, next, .. })) => {
        f(i, key)
        continue i + 1, next
      }
      (_, None) => break
    }
  }
}

///|
/// Clears the set, removing all keys. Keeps the allocated space.
pub fn[K] Set::clear(self : Set[K]) -> Unit {
  self.entries.fill(None)
  self.size = 0
  self.head = None
  self.tail = -1
}

///|
/// Returns the iterator of the hash set, provide elements in the order of insertion.
#alias(iterator, deprecated)
pub fn[K] Set::iter(self : Set[K]) -> Iter[K] {
  let mut curr_entry = self.head
  let len = self.size
  let mut remaining = len
  Iter::new(
    fn() {
      guard remaining > 0 && curr_entry is Some({ key, next, .. }) else { None }
      curr_entry = next
      remaining -= 1
      Some(key)
    },
    size_hint=len,
  )
}

///|
/// Converts the hash set to an array.
pub fn[K] Set::to_array(self : Set[K]) -> Array[K] {
  let arr = Array::new(capacity=self.size)
  for x = self.head {
    match x {
      Some({ key, next, .. }) => {
        arr.push(key)
        continue next
      }
      None => break
    }
  }
  arr
}

///|
pub impl[K : Hash + Eq] Eq for Set[K] with fn equal(self, other) {
  guard self.size == other.size else { return false }
  for k in self {
    guard other.contains(k) else { return false }
  } nobreak {
    true
  }
}

///|
/// Create from `iter`.
#as_free_fn
#alias(from_iterator, deprecated)
#as_free_fn(from_iterator, deprecated)
pub fn[K : Hash + Eq] Set::from_iter(iter : Iter[K]) -> Set[K] {
  let m = new_set(default_init_capacity)
  while iter.next() is Some(e) {
    m.add(e)
  }
  m
}

///|
pub impl[K] Default for Set[K] with fn default() {
  new_set(default_init_capacity)
}

///|
/// Copy the set, creating a new set with the same keys and order of insertion.
#alias(clone, deprecated)
pub fn[K] Set::copy(self : Set[K]) -> Set[K] {
  // copy structure
  let other = {
    capacity: self.capacity,
    entries: FixedArray::make(self.capacity, None),
    size: self.size,
    capacity_mask: self.capacity_mask,
    grow_at: self.grow_at,
    head: None,
    tail: self.tail,
  }
  if self.size == 0 {
    return other
  }
  guard! self.entries[self.tail] is Some(last)
  for entry = last, idx = self.tail, next = (None : Entry[K]?) {
    match (entry, idx, next) {
      ({ prev, psl, hash, key, .. }, idx, next) => {
        let new_entry = { prev, next, psl, hash, key }
        other.entries[idx] = Some(new_entry)
        if prev != -1 {
          continue self.entries[prev].unwrap(), prev, Some(new_entry)
        } else {
          other.head = Some(new_entry)
          break
        }
      }
    }
  }
  other
}

///|
/// Returns a new set containing elements in `self` that are not in `other`.
pub fn[K : Hash + Eq] Set::difference(self : Set[K], other : Set[K]) -> Set[K] {
  let m = new_set(default_init_capacity)
  self.each(k => if !other.contains(k) { m.add(k) })
  m
}

///|
/// Returns a new set containing elements in exactly one of the two sets.
pub fn[K : Hash + Eq] Set::symmetric_difference(
  self : Set[K],
  other : Set[K],
) -> Set[K] {
  let m = new_set(default_init_capacity)
  self.each(k => if !other.contains(k) { m.add(k) })
  other.each(k => if !self.contains(k) { m.add(k) })
  m
}

///|
/// Returns a new set containing all elements from both sets.
pub fn[K : Hash + Eq] Set::union(self : Set[K], other : Set[K]) -> Set[K] {
  let m = new_set(default_init_capacity)
  self.each(k => m.add(k))
  other.each(k => m.add(k))
  m
}

///|
/// Returns a new set containing only elements present in both sets.
pub fn[K : Hash + Eq] Set::intersection(
  self : Set[K],
  other : Set[K],
) -> Set[K] {
  let m = new_set(default_init_capacity)
  self.each(k => if other.contains(k) { m.add(k) })
  m
}

///|
pub impl[X : ToJson] ToJson for Set[X] with fn to_json(self) {
  let res = Array::new(capacity=self.size)
  for v in self {
    res.push(v.to_json())
  }
  Json::array(res)
}

///|
/// Returns `true` if the two sets have no elements in common.
pub fn[K : Hash + Eq] Set::is_disjoint(self : Set[K], other : Set[K]) -> Bool {
  if self.length() <= other.length() {
    for k in self {
      if other.contains(k) {
        return false
      }
    }
  } else {
    for k in other {
      if self.contains(k) {
        return false
      }
    }
  }
  true
}

///|
/// Returns `true` if every element of `self` is also in `other`.
pub fn[K : Hash + Eq] Set::is_subset(self : Set[K], other : Set[K]) -> Bool {
  if self.length() <= other.length() {
    for k in self {
      if !other.contains(k) {
        return false
      }
    }
    true
  } else {
    false
  }
}

///|
/// Returns `true` if every element of `other` is also in `self`.
pub fn[K : Hash + Eq] Set::is_superset(self : Set[K], other : Set[K]) -> Bool {
  other.is_subset(self)
}

///|
/// Intersection of two hash sets.
pub impl[K : Hash + Eq] BitAnd for Set[K] with fn land(self, other) {
  self.intersection(other)
}

///|
/// Union of two hash sets.
pub impl[K : Hash + Eq] BitOr for Set[K] with fn lor(self, other) {
  self.union(other)
}

///|
/// Symmetric difference of two hash sets.
pub impl[K : Hash + Eq] BitXOr for Set[K] with fn lxor(self, other) {
  self.symmetric_difference(other)
}

///|
/// Difference of two hash sets.
pub impl[K : Hash + Eq] Sub for Set[K] with fn sub(self, other) {
  self.difference(other)
}

///|