// 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.

// An implementation of HAMT (Hash Array Mapped Trie) in MoonBit.
//
// Hash-Array-Mapped-Trie (HAMT) is a persistent hash-table data structure.
// It is a trie over the hash of keys (i.e. strings of binary digits)
//
// Every level in a HAMT can have up to 32 branches (5 digits),
// so HAMT has a tree height of at most 7,
// and is more efficient compared to most other tree data structures.
//
// HAMT uses bitmap-based sparse array to avoid space waste
//
// Some references:
// - 
// - 

///|
priv struct BuildEntry[A] {
  value : A
  path : @path.Path // Full path; bulk construction tracks consumed segments by depth.
}

///|
let bulk_build_threshold = 64

///|
/// Create a new instance.
#as_free_fn
pub fn[A] HashSet::new() -> HashSet[A] {
  { data: None }
}

///|
/// Lookup a value from the hash set
pub fn[A : Eq + Hash] HashSet::contains(self : HashSet[A], key : A) -> Bool {
  self.data is Some(node) && node.contains(key, @path.of(key))
}

///|
fn[A : Eq] Node::contains(self : Node[A], key : A, path : @path.Path) -> Bool {
  for s = self, p = path {
    match (s, p) {
      (Leaf(key1, bucket), _) => break key == key1 || bucket.contains(key)
      (Flat(key1, path1), path) => break path == path1 && key == key1
      (Branch(children), path) => {
        let idx = path.idx()
        if children.get(idx) is Some(child) {
          continue child, path.next()
        }
        break false
      }
    }
  }
}

///|
/// require: key1 != key2, path1 and path2 has the same length
#owned(key1, key2)
fn[A] join_2(
  key1 : A,
  path1 : @path.Path,
  key2 : A,
  path2 : @path.Path,
) -> Node[A] {
  let idx1 = path1.idx()
  let idx2 = path2.idx()
  if idx1 == idx2 {
    let node = if path1.is_last() {
      Leaf(key2, @list.singleton(key1))
    } else {
      join_2(key1, path1.next(), key2, path2.next())
    }
    Branch(@sparse_array.singleton(idx1, node))
  } else {
    let (node1, node2) = if path1.is_last() {
      (Leaf(key1, @list.empty()), Leaf(key2, @list.empty()))
    } else {
      (Flat(key1, path1.next()), Flat(key2, path2.next()))
    }
    Branch(@sparse_array.doubleton(idx1, node1, idx2, node2))
  }
}

///|
fn[A : Eq] Node::add_with_path(
  self : Node[A],
  key : A,
  path : @path.Path,
) -> Node[A] {
  match self {
    Leaf(key1, bucket) =>
      if key == key1 || bucket.contains(key) {
        self
      } else {
        Leaf(key, bucket.add(key1))
      }
    Flat(key1, path1) =>
      if path == path1 && key == key1 {
        self
      } else {
        join_2(key1, path1, key, path)
      }
    Branch(children) => {
      let idx = path.idx()
      match children.get(idx) {
        Some(child) => {
          let child = child.add_with_path(key, path.next())
          Branch(children.replace(idx, child))
        }
        None => {
          let child = Flat(key, path.next())
          Branch(children.add(idx, child))
        }
      }
    }
  }
}

///|
fn[A : Eq] BuildEntry::add_to_node(
  self : BuildEntry[A],
  node : Node[A]?,
  depth : Int,
) -> Node[A] {
  let path = self.path.advance(depth)
  match node {
    None => Flat(self.value, path)
    Some(node) => node.add_with_path(self.value, path)
  }
}

///|
fn[A : Eq] build_hashset_node_by_add(
  entries : Array[BuildEntry[A]],
  start : Int,
  end : Int,
  depth : Int,
  reverse : Bool,
) -> Node[A] {
  if reverse {
    for i = end, node = (None : Node[A]?) {
      if i == start {
        break node.unwrap()
      }
      let node = entries[i - 1].add_to_node(node, depth)
      continue i - 1, Some(node)
    }
  } else {
    for i = start, node = (None : Node[A]?) {
      if i == end {
        break node.unwrap()
      }
      let node = entries[i].add_to_node(node, depth)
      continue i + 1, Some(node)
    }
  }
}

///|
fn[A : Eq] build_hashset_node_range(
  entries : Array[BuildEntry[A]],
  start : Int,
  end : Int,
  depth : Int,
  reverse : Bool,
) -> Node[A] {
  if end - start == 1 {
    let entry = entries[start]
    return Flat(entry.value, entry.path.advance(depth))
  }
  if depth == 5 {
    return build_hashset_node_by_add(entries, start, end, depth, reverse)
  }
  let starts = FixedArray::make(32, 0)
  for i in start.. Flat(value, path.push(first_idx))
      _ => Branch(@sparse_array.singleton(first_idx, first_child))
    }
  }
  let indices = FixedArray::make(child_count, first_idx)
  let children = FixedArray::make(child_count, first_child)
  for idx = first_idx + 1, out = 1; idx < 32; {
    if nexts[idx] == starts[idx] {
      continue idx + 1, out
    }
    indices[out] = idx
    children[out] = build_hashset_node_range(
      partitioned,
      starts[idx],
      nexts[idx],
      depth + 1,
      reverse,
    )
    continue idx + 1, out + 1
  }
  Branch(@sparse_array.from_sorted_fixed_array(indices, children))
}

///|
fn[A : Eq + Hash] hash_set_from_iter_by_add(iter : Iter[A]) -> HashSet[A] {
  iter.fold(init=new(), (s, e) => s.add(e))
}

///|
fn[A : Eq + Hash] hash_set_from_array_by_add(arr : ArrayView[A]) -> HashSet[A] {
  for n = arr.length(), set = new() {
    match (n, set) {
      (0, set) => break set
      (n, set) => {
        let k = arr[n - 1]
        continue n - 1, set.add(k)
      }
    }
  }
}

///|
fn[A : Eq + Hash] hash_set_from_array(arr : ArrayView[A]) -> HashSet[A] {
  if arr.length() <= bulk_build_threshold {
    return hash_set_from_array_by_add(arr)
  }
  let entries = Array::makei(arr.length(), i => {
    let value = arr[i]
    { value, path: @path.of(value) }
  })
  {
    data: Some(build_hashset_node_range(entries, 0, entries.length(), 0, true)),
  }
}

///|
/// Add a key to the hashset.
pub fn[A : Eq + Hash] HashSet::add(self : HashSet[A], key : A) -> HashSet[A] {
  {
    data: match self.data {
      None => Some(Flat(key, @path.of(key)))
      Some(node) => Some(node.add_with_path(key, @path.of(key)))
    },
  }
}

///|
/// Remove an element from a set
pub fn[A : Eq + Hash] HashSet::remove(self : HashSet[A], key : A) -> HashSet[A] {
  {
    data: match self.data {
      None => None
      Some(node) => node.remove_with_path(key, @path.of(key))
    },
  }
}

///|
fn[A : Eq] Node::remove_with_path(
  self : Node[A],
  key : A,
  path : @path.Path,
) -> Node[A]? {
  match self {
    Leaf(key1, bucket) =>
      if key1 == key {
        match bucket {
          Empty => None
          More(key2, tail=xs) => Some(Leaf(key2, xs))
        }
      } else if bucket.find_index(x => key.equal(x)) is Some(index) {
        Some(Leaf(key1, bucket.remove_at(index)))
      } else {
        Some(self)
      }
    Flat(key1, path1) =>
      if path == path1 && key == key1 {
        None
      } else {
        Some(self)
      }
    Branch(children) => {
      let idx = path.idx()
      match children.get(idx) {
        None => Some(self)
        Some(child) => {
          let new_child = child.remove_with_path(key, path.next())
          let new_children = match (children.length(), new_child) {
            (1, None) => return None
            (_, None) => children.remove(idx)
            (_, Some(new_child)) => children.replace(idx, new_child)
          }
          match new_children.data {
            [Flat(key1, path1)] =>
              Some(Flat(key1, path1.push(new_children.elem_info.first_idx())))
            _ => Some(Branch(new_children))
          }
        }
      }
    }
  }
}

///|
/// Calculate the size of a set.
///
/// WARNING: this operation is `O(N)` in set size
#alias(size, deprecated)
pub fn[A] HashSet::length(self : HashSet[A]) -> Int {
  fn node_size(node) {
    match node {
      Leaf(_, bucket) => 1 + bucket.length()
      Flat(_) => 1
      Branch(children) =>
        for child in children.data; total_size = 0 {
          continue total_size + node_size(child)
        } nobreak {
          total_size
        }
    }
  }

  match self.data {
    None => 0
    Some(node) => node_size(node)
  }
}

///|
/// Union two hashsets
pub fn[K : Eq] HashSet::union(
  self : HashSet[K],
  other : HashSet[K],
) -> HashSet[K] {
  fn go(node1, node2) {
    match (node1, node2) {
      (node, Flat(key, path)) | (Flat(key, path), node) =>
        node.add_with_path(key, path)
      (Branch(children1), Branch(children2)) =>
        Branch(children1.union(children2, go))
      (Leaf(key1, bucket1), Leaf(key2, bucket2)) => {
        let keys1 = bucket1.add(key1)
        let keys2 = bucket2.add(key2)
        match keys1.filter(k => !keys2.contains(k)) {
          Empty => node2
          More(head, tail~) => Leaf(key2, bucket2 + tail.add(head))
        }
      }
      _ => abort("Unreachable")
    }
  }

  {
    data: match (self.data, other.data) {
      (None, x) | (x, None) => x
      (Some(a), Some(b)) => Some(go(a, b))
    },
  }
}

///|
/// Intersect two hashsets
pub fn[K : Eq] HashSet::intersection(
  self : HashSet[K],
  other : HashSet[K],
) -> HashSet[K] {
  fn go(node1, node2) {
    match (node1, node2) {
      (node, Flat(key, path) as flat) | (Flat(key, path) as flat, node) =>
        if node.contains(key, path) {
          Some(flat)
        } else {
          None
        }
      (Branch(children1), Branch(children2)) =>
        match children1.intersection(children2, go) {
          None => None
          Some({ data: [Flat(key, path)], elem_info }) =>
            Some(Flat(key, path.push(elem_info.first_idx())))
          Some(children) => Some(Branch(children))
        }
      (Leaf(key1, bucket1), Leaf(key2, bucket2)) => {
        let keys1 = bucket1.add(key1)
        let keys2 = bucket2.add(key2)
        match keys1.filter(k => keys2.contains(k)) {
          Empty => None
          More(head, tail~) => Some(Leaf(head, tail))
        }
      }
      _ => abort("Unreachable")
    }
  }

  {
    data: match (self.data, other.data) {
      (None, _) | (_, None) => None
      (Some(a), Some(b)) => go(a, b)
    },
  }
}

///|
/// Difference of two hashsets: elements in `self` but not in `other`
pub fn[K : Eq] HashSet::difference(
  self : HashSet[K],
  other : HashSet[K],
) -> HashSet[K] {
  fn go(node1 : Node[_], node2) {
    match (node1, node2) {
      (node, Flat(k, path)) => node.remove_with_path(k, path)
      (Flat(key, path) as flat, node) =>
        if node.contains(key, path) {
          None
        } else {
          Some(flat)
        }
      (Branch(children1), Branch(children2)) =>
        match children1.difference(children2, go) {
          None => None
          Some({ data: [Flat(key, path)], elem_info }) =>
            Some(Flat(key, path.push(elem_info.first_idx())))
          Some(children) => Some(Branch(children))
        }
      (Leaf(key1, bucket1), Leaf(key2, bucket2)) => {
        let keys1 = bucket1.add(key1)
        let keys2 = bucket2.add(key2)
        match keys1.filter(k => !keys2.contains(k)) {
          Empty => None
          More(head, tail~) => Some(Leaf(head, tail))
        }
      }
      _ => abort("Unreachable")
    }
  }

  match (self.data, other.data) {
    (None, _) => { data: None }
    (_, None) => self
    (Some(a), Some(b)) => { data: go(a, b) }
  }
}

///|
/// Returns true if the hash set is empty.
pub fn[A] HashSet::is_empty(self : HashSet[A]) -> Bool {
  self.data is None
}

///|
/// Iterate through the elements in a hash set
pub fn[A] HashSet::each(
  self : HashSet[A],
  f : (A) -> Unit raise?,
) -> Unit raise? {
  fn go(node) raise? {
    match node {
      Leaf(k, bucket) => {
        f(k)
        bucket.each(f)
      }
      Flat(k, _) => f(k)
      Branch(children) => children.each(go)
    }
  }

  if self.data is Some(node) {
    go(node)
  }
}

///|
priv enum CurrNode[A] {
  Tree(Node[A])
  Bucket(@list.List[A])
}

///|
/// Converted to Iter
#alias(iterator, deprecated)
pub fn[A] HashSet::iter(self : HashSet[A]) -> Iter[A] {
  let empty = Bucket(@list.new())
  let mut curr_node = match self.data {
    Some(tree) => Tree(tree)
    None => empty
  }
  let mut curr_index = 0
  let parents = []
  Iter::new(fn() {
    for cn = curr_node {
      match cn {
        Tree(Flat(x, _)) => {
          curr_node = empty
          break Some(x)
        }
        Tree(Leaf(x, bucket)) => {
          curr_node = Bucket(bucket)
          break Some(x)
        }
        Bucket(More(x, tail~)) => {
          curr_node = Bucket(tail)
          break Some(x)
        }
        Tree(Branch(children)) as n if curr_index < children.length() => {
          let child = children.data[curr_index]
          parents.push((n, curr_index + 1))
          curr_index = 0
          continue Tree(child)
        }
        Bucket(Empty) | Tree(Branch(_)) if parents.pop()
          is Some((parent, parent_index)) => {
          curr_node = parent
          curr_index = parent_index
          continue parent
        }
        Bucket(Empty) | Tree(Branch(_)) => break None
      }
    }
  })
}

///|
/// Creates a hash set from an iterator of values.
#as_free_fn
#alias(from_iterator, deprecated)
#as_free_fn(from_iterator, deprecated)
pub fn[A : Eq + Hash] HashSet::from_iter(iter : Iter[A]) -> HashSet[A] {
  if iter.size_hint() is Some(len) && len <= bulk_build_threshold {
    return hash_set_from_iter_by_add(iter)
  }
  let entries = match iter.size_hint() {
    Some(len) => Array::new(capacity=len)
    None => []
  }
  iter.each(value => entries.push({ value, path: @path.of(value) }))
  if entries.is_empty() {
    new()
  } else {
    {
      data: Some(
        build_hashset_node_range(entries, 0, entries.length(), 0, false),
      ),
    }
  }
}

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

///|
pub impl[A : Show] Show for HashSet[A] with fn output(self, logger) {
  logger.write_iter(
    self.iter(),
    prefix="@immut/hashset.from_array([",
    suffix="])",
  )
}

///|
/// Creates a hash set from an array of values.
#as_free_fn(deprecated="Use @immut/hashset.HashSet([...]) instead")
#alias(of, deprecated="Use @immut/hashset.HashSet([...]) instead")
#as_free_fn(of, deprecated="Use @immut/hashset.HashSet([...]) instead")
#deprecated("Use @immut/hashset.HashSet([...]) instead")
pub fn[A : Eq + Hash] HashSet::from_array(arr : ArrayView[A]) -> HashSet[A] {
  hash_set_from_array(arr)
}

///|
/// Creates a hash set from an array of values.
///
/// # Example
///
/// ```mbt check
/// test {
///   let set = @hashset.HashSet([3, 1, 2, 3])
///   @test.assert_eq(set.contains(1), true)
///   @test.assert_eq(set.contains(4), false)
/// }
/// ```
pub fn[A : Eq + Hash] HashSet::HashSet(arr : ArrayView[A]) -> HashSet[A] {
  hash_set_from_array(arr)
}

///|
pub impl[A : Hash] Hash for HashSet[A] with fn hash_combine(self, hasher) {
  hasher.combine(self.iter().fold(init=0, (x, y) => x ^ y.hash()))
}

///|
impl[A : Eq] Eq for Node[A] with fn equal(self, other) {
  if physical_equal(self, other) {
    return true
  }
  match (self, other) {
    (Leaf(x, xs), Leaf(y, ys)) =>
      xs.length() == ys.length() &&
      ({
        let keys1 = xs.add(x)
        let keys2 = ys.add(y)
        keys1.iter().all(k => keys2.contains(k))
      })
    (Flat(x, pathx), Flat(y, pathy)) => pathx == pathy && x == y
    (Branch(xs), Branch(ys)) => xs == ys
    _ => false
  }
}

///|
test "hash" {
  @test.assert_eq(
    Hash::hash(HashSet([1, 2, 3, 4])),
    Hash::hash(HashSet([3, 2]).add(1).add(4)),
  )
  @test.assert_not_eq(
    Hash::hash(HashSet([1, 2, 3])),
    Hash::hash(HashSet([1, 2, 4])),
  )
}

///|
test "eq" {
  @test.assert_eq(HashSet([1, 2, 3, 4]), HashSet([3, 2]).add(1).add(4))
  @test.assert_not_eq(HashSet([1, 2, 3]), HashSet([1, 2, 4]))
}