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

///|
/// Create an empty map.
#as_free_fn
#alias(empty, deprecated)
pub fn[K, V] SortedMap::new() -> SortedMap[K, V] {
  Empty
}

///|
/// Create a map with a single key-value pair.
#as_free_fn
#owned(key, value)
pub fn[K, V] SortedMap::singleton(key : K, value : V) -> SortedMap[K, V] {
  Tree(key, value~, size=1, Empty, Empty)
}

///|
/// Check if the map contains a key.
/// O(log n).
pub fn[K : Compare, V] SortedMap::contains(
  self : SortedMap[K, V],
  key : K,
) -> Bool {
  for x = self {
    match x {
      Empty => break false
      Tree(k, l, r, ..) => {
        let c = key.compare(k)
        if c == 0 {
          break true
        } else if c < 0 {
          continue l
        } else {
          continue r
        }
      }
    }
  }
}

///|
/// Get the number of key-value pairs in the map.
#alias(size, deprecated)
#inline
pub fn[K, V] SortedMap::length(self : SortedMap[K, V]) -> Int {
  match self {
    Empty => 0
    Tree(_) as t => t.size
  }
}

///|
/// Return whether the value empty.
pub fn[K, V] SortedMap::is_empty(self : SortedMap[K, V]) -> Bool {
  self.length() == 0
}

///|
#inline
#owned(key, value, l, r)
fn[K, V] make_tree(
  key : K,
  value : V,
  l : SortedMap[K, V],
  r : SortedMap[K, V],
) -> SortedMap[K, V] {
  // length() is #inline, so the size match collapses into make_tree, which
  // is itself #inline into balance / merge / add.
  let size = l.length() + r.length() + 1
  Tree(key, value~, size~, l, r)
}

///|
/// Get the value associated with a key.
/// O(log n).
#alias(lookup, deprecated)
pub fn[K : Compare, V] SortedMap::get(self : SortedMap[K, V], key : K) -> V? {
  for x = self {
    match x {
      Empty => break None
      Tree(k, value~, l, r, ..) => {
        let c = key.compare(k)
        if c == 0 {
          break Some(value)
        } else if c < 0 {
          continue l
        } else {
          continue r
        }
      }
    }
  }
}

///|
/// Get the value associated with a key.
/// O(log n).
#alias("_[_]")
pub fn[K : Compare, V] SortedMap::at(self : SortedMap[K, V], key : K) -> V {
  for x = self {
    match x {
      Empty => break panic()
      Tree(k, value~, l, r, ..) => {
        let c = key.compare(k)
        if c == 0 {
          break value
        } else if c < 0 {
          continue l
        } else {
          continue r
        }
      }
    }
  }
}

///|
/// Iterate over the key-value pairs in the map.
pub fn[K, V] SortedMap::each(
  self : SortedMap[K, V],
  f : (K, V) -> Unit,
) -> Unit {
  match self {
    Empty => ()
    Tree(k, value~, l, r, ..) => {
      l.each(f)
      f(k, value)
      r.each(f)
    }
  }
}

///|
/// Iterate over the key-value pairs with index.
pub fn[K, V] SortedMap::eachi(
  self : SortedMap[K, V],
  f : (Int, K, V) -> Unit,
) -> Unit {
  fn do_eachi(m : SortedMap[K, V], f, i) {
    match m {
      Empty => ()
      Tree(k, value~, l, r, ..) => {
        do_eachi(l, f, i)
        f(l.length() + i, k, value)
        do_eachi(r, f, l.length() + i + 1)
      }
    }
  }

  do_eachi(self, f, 0)
}

///|
/// Maps over the key-value pairs in the map.
#alias(map_with_key, deprecated)
pub fn[K, X, Y] SortedMap::map(
  self : SortedMap[K, X],
  f : (K, X) -> Y,
) -> SortedMap[K, Y] {
  match self {
    Empty => Empty
    Tree(k, value~, l, r, size~) =>
      Tree(k, value=f(k, value), size~, l.map(f), r.map(f))
  }
}

///|
/// Post-order fold.
/// O(n).
#alias(foldr_with_key)
pub fn[K, V, A] SortedMap::rev_fold(
  self : SortedMap[K, V],
  f : (A, K, V) -> A,
  init~ : A,
) -> A {
  fn go(m : SortedMap[K, V], acc) {
    match m {
      Empty => acc
      Tree(k, value~, l, r, ..) => go(l, f(go(r, acc), k, value))
    }
  }

  go(self, init)
}

///|
/// Pre-order fold.
/// O(n).
#alias(foldl_with_key, deprecated)
pub fn[K, V, A] SortedMap::fold(
  self : SortedMap[K, V],
  f : (A, K, V) -> A,
  init~ : A,
) -> A {
  fn go(m : SortedMap[K, V], acc) {
    match m {
      Empty => acc
      Tree(k, value~, l, r, ..) => go(r, f(go(l, acc), k, value))
    }
  }

  go(self, init)
}

///|
fn[K : Show, V : Show] SortedMap::debug_tree(self : SortedMap[K, V]) -> String {
  match self {
    Empty => "_"
    Tree(k, value~, l, r, ..) => {
      let l = l.debug_tree()
      let r = r.debug_tree()
      "(\{k},\{value},\{l},\{r})"
    }
  }
}

///|
/// Build a map from an array of key-value pairs.
/// O(n*log n).
#as_free_fn(deprecated="Use @immut/sorted_map.SortedMap([...]) instead")
#alias(of, deprecated="Use @immut/sorted_map.SortedMap([...]) instead")
#as_free_fn(of, deprecated="Use @immut/sorted_map.SortedMap([...]) instead")
#deprecated("Use @immut/sorted_map.SortedMap([...]) instead")
pub fn[K : Compare, V] SortedMap::from_array(
  array : ArrayView[(K, V)],
) -> SortedMap[K, V] {
  for item in array; mp = Empty {
    let (k, v) = item
    continue mp.add(k, v)
  } nobreak {
    mp
  }
}

///|
/// Build a map from an array of key-value pairs.
/// O(n*log n).
///
/// # Example
///
/// ```mbt check
/// test {
///   let m = @sorted_map.SortedMap([(3, "c"), (1, "a"), (2, "b")])
///   @test.assert_eq(m.get(1), Some("a"))
///   @test.assert_eq(m.get(3), Some("c"))
/// }
/// ```
pub fn[K : Compare, V] SortedMap::SortedMap(
  array : ArrayView[(K, V)],
) -> SortedMap[K, V] {
  for item in array; mp = Empty {
    let (k, v) = item
    continue mp.add(k, v)
  } nobreak {
    mp
  }
}

///|
/// Returns an iterator over key-value pairs in ascending key order.
#alias(iterator, deprecated)
pub fn[K, V] SortedMap::iter(self : SortedMap[K, V]) -> Iter[(K, V)] {
  let mut curr_node = self
  let parents = []
  Iter::new(
    fn() {
      for x = curr_node {
        match x {
          Tree(k, value~, Empty, r, ..) => {
            curr_node = r
            break Some((k, value))
          }
          Tree(k, value~, l, r, ..) => {
            parents.push((k, value, r))
            continue l
          }
          Empty if parents.pop() is Some((k, v, r)) => {
            curr_node = r
            break Some((k, v))
          }
          Empty => break None
        }
      }
    },
    size_hint=self.length(),
  )
}

///|
/// Returns a two-element iterator over key-value pairs in ascending key order.
#alias(iterator2, deprecated)
pub fn[K, V] SortedMap::iter2(self : SortedMap[K, V]) -> Iter2[K, V] {
  self.iter()
}

///|
/// Returns an iterator over all key-value pairs in the map where keys fall
/// within the inclusive range `[low, high]`.
///
/// The iterator yields key-value pairs in ascending key order. Keys equal to
/// `low` or `high` are included if present.
///
/// # Arguments
///
/// * `low` - The lower bound of the range (inclusive).
/// * `high` - The upper bound of the range (inclusive).
///
/// # Returns
///
/// An `Iter2[K, V]` that yields key-value pairs `(key, value)` where
/// `low <= key <= high`.
///
/// # Performance
///
/// Time complexity is O(log n + k) where n is the size of the map and k is the
/// number of elements in the range. The algorithm efficiently prunes subtrees
/// that fall entirely outside the range.
///
/// # Behavior
///
/// * If `low > high`, the iterator yields no elements.
/// * If the range contains no keys from the map, the iterator yields no elements.
/// * The iterator is single-use; create a new one for multiple traversals.
///
/// # Example
///
/// ```mbt check
/// test {
///   let map = @sorted_map.SortedMap([
///     (1, "a"),
///     (2, "b"),
///     (3, "c"),
///     (4, "d"),
///     (5, "e"),
///   ])
///   let result = []
///   for k, v in map.range(low=2, high=4) {
///     result.push((k, v))
///   }
///   @test.assert_eq(result, [(2, "b"), (3, "c"), (4, "d")])
/// }
/// ```
pub fn[K : Compare, V] SortedMap::range(
  self : SortedMap[K, V],
  low~ : K,
  high~ : K,
) -> Iter2[K, V] {
  let todo_list = []
  let mut next_node = self
  Iter2::new(fn() {
    for x = next_node {
      match x {
        Tree(k, value~, l, r, ..) => {
          let cmp_key_low = k.compare(low)
          let cmp_key_high = k.compare(high)
          if cmp_key_low < 0 {
            // k < low, skip left subtree and current node, visit right
            continue r
          } else if cmp_key_high > 0 {
            // k > high, skip right subtree and current node, visit left
            continue l
          } else if l is Empty {
            next_node = r
            break Some((k, value))
          } else {
            todo_list.push((k, value, r))
            continue l
          }
        }
        Empty if todo_list.pop() is Some((k, value, r)) => {
          next_node = r
          break Some((k, value))
        }
        Empty => break None
      }
    }
  })
}

///|
/// Creates a sorted map from an iterator of key-value pairs.
#as_free_fn
#alias(from_iterator, deprecated)
#as_free_fn(from_iterator, deprecated)
pub fn[K : Compare, V] SortedMap::from_iter(
  iter : Iter[(K, V)],
) -> SortedMap[K, V] {
  iter.fold(init=new(), (m, e) => m.add(e.0, e.1))
}

///|
/// Return all keys of the map in ascending order.
pub fn[K, V] SortedMap::keys_as_iter(self : SortedMap[K, V]) -> Iter[K] {
  self.iter().map(p => p.0)
}

///|
/// Return all elements of the map in the ascending order of their keys.
pub fn[K, V] SortedMap::values(self : SortedMap[K, V]) -> Iter[V] {
  self.iter().map(p => p.1)
}

///|
/// Return all keys of the map in descending order.
///
/// # Example
///
/// ```mbt check
/// test {
///   let map = @sorted_map.SortedMap([("a", 1), ("b", 2), ("c", 3)])
///   let keys = map.rev_keys().collect()
///   @test.assert_eq(keys, ["c", "b", "a"])
/// }
/// ```
pub fn[K, V] SortedMap::rev_keys(self : SortedMap[K, V]) -> Iter[K] {
  let mut curr_node = self
  let parents = []
  Iter::new(
    fn() {
      for x = curr_node {
        match x {
          Tree(k, l, Empty, ..) => {
            curr_node = l
            break Some(k)
          }
          Tree(k, l, r, ..) => {
            parents.push((k, l))
            continue r
          }
          Empty if parents.pop() is Some((k, l)) => {
            curr_node = l
            break Some(k)
          }
          Empty => break None
        }
      }
    },
    size_hint=self.length(),
  )
}

///|
/// Return all values of the map in descending order of their keys.
///
/// # Example
///
/// ```mbt nocheck
///   let map = @sorted_map.SortedMap([("a", 1), ("b", 2), ("c", 3)])
///   let values = map.rev_values().collect()
///   @test.assert_eq(values, [3, 2, 1])
/// ```
pub fn[K, V] SortedMap::rev_values(self : SortedMap[K, V]) -> Iter[V] {
  let mut curr_node = self
  let parents = []
  Iter::new(
    fn() {
      for x = curr_node {
        match x {
          Tree(_k, value~, l, Empty, ..) => {
            curr_node = l
            break Some(value)
          }
          Tree(_k, value~, l, r, ..) => {
            parents.push((value, l))
            continue r
          }
          Empty if parents.pop() is Some((v, l)) => {
            curr_node = l
            break Some(v)
          }
          Empty => break None
        }
      }
    },
    size_hint=self.length(),
  )
}

///|
/// Convert to `json`.
pub fn[K : Show, V : ToJson] SortedMap::to_json(self : SortedMap[K, V]) -> Json {
  ToJson::to_json(self)
}