fractional_indexing

//// This library implements fractional indexing outlined in these two blog posts:
////  * [Real-time Editing of Ordered Sequences](https://www.figma.com/blog/realtime-editing-of-ordered-sequences/)
////  * [Implementing Fractional Indexing](https://observablehq.com/@dgreensp/implementing-fractional-indexing)
////
//// Fractional indexing is useful in maintaining an ordered sequence of keys where you want to insert keys at
//// arbitary locations in the list.  The original Figma post above outlined a system used in their product to
//// achieve eventual consistency using a decimal index.  For example, in a list with keys `[0, 1]` you could add
//// a new item between them with index `0.5`.
////
//// Indexing using floats runs into precision issues in Javascript and starts to yield very long keys.  This library
//// implements a system of variable length integer encoding followed by fractional indexing to keep keys short.  Instead of
//// sorting items by their numerical index, these keys sort lexicographically using the algorithm at the second link above.
////
//// ## Example
////
//// ```gleam
//// import fractional_indexing as fi
////
//// // start with a list of items
//// let items = [1, 3]
////
//// // create some keys for them, here we will use the bulk API
//// // to generate several keys at once and zip them up
//// let keys = fi.n_first(num_keys: items |> list.length)
////
//// let zipped = list.zip(keys, items)
//// // [#("a0", 1), #("a1", 3)]
////
//// // now you want to add an item between these two values, so
//// // generate a key to index between `a0` and `a1`. Note
//// // that key generation can fail in exotic situations, but under
//// // normal use should return a suitable key
//// let assert [key1, key2] = keys // a0, a1
//// let assert Ok(key_between) = fi.between(key1, key2)
//// let zipped = [#(key_between, 2), ..zipped]
////
//// // sort them to get them in the order you want using string comparison
//// // for the keys
//// echo zipped |> list.sort(by: fn(a, b) { string.compare(a.0, b.0) })
//// // [#("a0", 1), #("a0V", 2) #("a1", 3)]
////
//// ```

import gleam/bool
import gleam/dict
import gleam/int
import gleam/list
import gleam/option.{type Option, None, Some}
import gleam/order
import gleam/result
import gleam/string

const base62 = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"

const integer_zero = "a0"

const smallest_integer = "A00000000000000000000000000"

/// Key generation is not guaranteed under all conditions. Billions of values are possible before exhausting the key space.  Keys
/// are encoded using a combination of variable length integer encoding and fractional indices.  This means that you cannot
/// generate a key starting from an arbitrary string.
///
pub type KeyError {
  /// When generating a key between two values, the first value must be lexicographically smaller than the first
  ErrWrongOrder
  /// Something about the encoding of the key is invalid or the keyspace is exhausted. You can generate billions of keys
  /// before running out.
  ErrInvalid
}

/// Generate a key that will sort later than the given key
///
pub fn after(key: String) -> Result(String, KeyError) {
  generate_key_between(key |> string.to_option, None, base62)
}

/// Generate a key that will sort before the given key
///
pub fn before(key: String) -> Result(String, KeyError) {
  generate_key_between(None, key |> string.to_option, base62)
}

/// Generate a key that will sort between key1 and key2
///
pub fn between(key1: String, key2: String) -> Result(String, KeyError) {
  generate_key_between(
    key1 |> string.to_option,
    key2 |> string.to_option,
    base62,
  )
}

/// Return the first key, defined as the zero-key `a0`.  This is useful when assigning a key to the first element of a list. Note that you cannot start a list of keys from any arbitrary string because the encoding
/// of the key matters.
///
pub fn first() -> String {
  let assert Ok(first) = generate_key_between(None, None, base62)
  first
}

/// Generate N consecutive keys starting from the zero key `a0`
///
pub fn n_first(num_keys n: Int) -> List(String) {
  let assert Ok(start) = before(first())
  let assert Ok(n_first) = n_between(start, "", n)
  n_first
}

/// Generate N keys after the given key
///
pub fn n_after(key: String, num_keys n: Int) -> Result(List(String), KeyError) {
  n_between(key, "", n)
}

/// Generate N keys before the given key
///
pub fn n_before(key: String, num_keys n: Int) -> Result(List(String), KeyError) {
  n_between("", key, n)
}

/// Generate N keys between key1 and key 2
///
/// ## Example
///
/// ```gleam
/// import fractional_indexing as fi
///
/// // start with a list of some items which we want to index between existing keys `a0` and `a1`
/// let items = [1, 2, 3]
///
/// // generate some keys
/// let assert Ok(keys) = fi.n_between("a0", "a1")
///
/// // do something with the keys
/// items
/// |> list.map2(keys, fn(item, key) {
///      #(key, item)
/// })
/// // output: [#("a0V", 1), #("a0k", 2), #("a0s", 3)]
/// ```
///
pub fn n_between(
  key1: String,
  key2: String,
  num_keys n: Int,
) -> Result(List(String), KeyError) {
  list.range(0, n - 1)
  |> list.try_fold([], fn(res, _) {
    case res {
      [] -> between(key1, key2) |> result.map(fn(key) { [key] })
      a -> {
        let assert Ok(last) = a |> list.last
        between(last, key2) |> result.map(fn(key) { list.append(a, [key]) })
      }
    }
  })
}

// generate the key using variable length integer encoding coupled with fractional index
// implements the algorithm at: https://observablehq.com/@dgreensp/implementing-fractional-indexing
@internal
pub fn generate_key_between(
  a: Option(String),
  b: Option(String),
  encoding: String,
) -> Result(String, KeyError) {
  // checks for properly encoded key and that the encoded key has a valid structure

  let identity = fn(v) { v }
  let check_ok = case a, b {
    Some(a), Some(b) ->
      [
        order_ok(a, b),
        order_key_ok(a),
        order_key_ok(b),
        encoding_ok(a <> b, encoding),
      ]
      |> list.try_each(identity)
    Some(a), None ->
      [order_key_ok(a), encoding_ok(a, encoding)] |> list.try_each(identity)
    None, Some(b) ->
      [order_key_ok(b), encoding_ok(b, encoding)] |> list.try_each(identity)
    None, None -> Ok(Nil)
  }
  case check_ok, a, b {
    Error(e), _, _ -> Error(e)
    Ok(Nil), a, b ->
      do_generate_key_between(a, b, encoding) |> option.to_result(ErrInvalid)
  }
}

fn do_generate_key_between(
  a: Option(String),
  b: Option(String),
  encoding: String,
) -> Option(String) {
  case a, b {
    None, None -> Some(integer_zero)
    None, Some(b) -> {
      let assert Some(int_part) = get_integer_part(b)
      let frac_part = b |> string.drop_start(int_part |> string.length)
      case int_part {
        ip if ip == smallest_integer ->
          Some(ip <> midpoint(None, frac_part |> string.to_option, encoding))
        ip -> {
          case string.compare(ip, b) == order.Lt {
            True -> Some(ip)
            _ -> decrement_integer(ip, encoding)
          }
        }
      }
    }
    Some(a), None -> {
      let assert Some(int_part) = get_integer_part(a)
      let frac_part = a |> string.drop_start(int_part |> string.length)
      case increment_integer(int_part, encoding) {
        Some(next) -> Some(next)
        None ->
          Some(
            int_part <> midpoint(frac_part |> string.to_option, None, encoding),
          )
      }
    }
    Some(a), Some(b) -> {
      let assert Some(int_part_a) = get_integer_part(a)
      let frac_part_a = a |> string.drop_start(int_part_a |> string.length)
      let assert Some(int_part_b) = get_integer_part(b)
      let frac_part_b = b |> string.drop_start(int_part_b |> string.length)
      case int_part_a == int_part_b {
        True ->
          Some(
            int_part_a
            <> midpoint(
              frac_part_a |> string.to_option,
              frac_part_b |> string.to_option,
              encoding,
            ),
          )
        False -> {
          let next_int_a = increment_integer(int_part_a, encoding)
          case next_int_a {
            Some(ia) -> {
              case string.compare(ia, b) == order.Lt {
                True -> next_int_a
                False ->
                  Some(
                    int_part_a
                    <> midpoint(frac_part_a |> string.to_option, None, encoding),
                  )
              }
            }
            None -> None
          }
        }
      }
    }
  }
}

// -----------------------------------------------------------------------------------
// VARIABLE LENGTH INTEGER ENCODING
// -----------------------------------------------------------------------------------

@internal
pub fn increment_integer(i: String, encoding: String) -> Option(String) {
  validate_integer_ok(i)
  |> option.then(fn(i) { do_increment_integer(i, encoding) })
}

@internal
pub fn decrement_integer(i: String, encoding: String) -> Option(String) {
  validate_integer_ok(i)
  |> option.then(fn(i) { do_decrement_integer(i, encoding) })
}

// increment the variable length encoded integer
fn do_increment_integer(i: String, encoding: String) -> Option(String) {
  let lut =
    list.zip(
      encoding |> string.split(""),
      list.range(0, encoding |> string.length),
    )
    |> dict.from_list

  let encoding_length = encoding |> string.length

  let assert [head, ..digits] = i |> string.split("")
  let #(carry, updated_reversed_digits) =
    digits
    |> list.reverse
    |> list.map_fold(True, fn(carry, c) {
      let assert Ok(c_index) = lut |> dict.get(c)
      let c_next = c_index + 1
      case carry, c_next {
        True, index if index == encoding_length -> #(True, "0")
        True, index -> #(
          False,
          encoding |> string.slice(at_index: index, length: 1),
        )
        False, _ -> #(False, c)
      }
    })
  let new_digits = updated_reversed_digits |> list.reverse |> string.join("")
  case carry, head {
    True, "Z" -> Some(integer_zero)
    True, "z" -> None
    True, head -> {
      let next_codepoint = get_codepoint(head) + 1
      let next_head = codepoint_to_string(next_codepoint)
      case next_codepoint {
        cp if cp > lower_a_charcode -> Some(next_head <> new_digits <> "0")
        _ -> {
          let pop_digits = new_digits |> string.drop_end(up_to: 1)
          Some(next_head <> pop_digits)
        }
      }
    }
    False, head -> Some(head <> new_digits)
  }
}

// decrement the variable length encoded integer
fn do_decrement_integer(i: String, encoding: String) -> Option(String) {
  let lut =
    list.zip(
      encoding |> string.split(""),
      list.range(0, encoding |> string.length),
    )
    |> dict.from_list

  let assert Ok(last_encoding_char) = encoding |> string.last

  let assert [head, ..digits] = i |> string.split("")
  let #(borrow, updated_reversed_digits) =
    digits
    |> list.reverse
    |> list.map_fold(True, fn(borrow, c) {
      let assert Ok(c_index) = lut |> dict.get(c)
      let c_next = c_index - 1
      case borrow, c_next {
        True, index if index == -1 -> #(True, last_encoding_char)
        True, index -> #(
          False,
          encoding |> string.slice(at_index: index, length: 1),
        )
        False, _ -> #(False, c)
      }
    })
  let new_digits = updated_reversed_digits |> list.reverse |> string.join("")
  case borrow, head {
    True, "a" -> Some("Z" <> last_encoding_char)
    True, "A" -> None
    True, head -> {
      let next_codepoint = get_codepoint(head) - 1
      let next_head = codepoint_to_string(next_codepoint)
      case next_codepoint {
        cp if cp < upper_z_charcode ->
          Some(next_head <> new_digits <> last_encoding_char)
        _ -> {
          let pop_digits = new_digits |> string.drop_end(up_to: 1)
          Some(next_head <> pop_digits)
        }
      }
    }
    False, head -> Some(head <> new_digits)
  }
}

fn get_codepoint(s: String) -> Int {
  let assert Ok(codepoint) = s |> string.to_utf_codepoints |> list.first
  codepoint |> string.utf_codepoint_to_int
}

fn codepoint_to_string(i: Int) -> String {
  let assert Ok(codepoint) = string.utf_codepoint(i)
  string.from_utf_codepoints([codepoint])
}

// useful char codes for the algorithm
const upper_a_charcode = 65

const upper_z_charcode = 90

const lower_a_charcode = 97

const lower_z_charcode = 122

fn get_integer_part(key: String) -> Option(String) {
  let first = key |> string.first |> result.unwrap("") |> string.to_option
  let key_length = key |> string.length
  case get_integer_length(first) {
    Some(len) if len <= key_length ->
      Some(key |> string.slice(at_index: 0, length: len))
    _ -> None
  }
}

fn get_integer_length(first: Option(String)) -> Option(Int) {
  case first {
    None -> None
    Some(first) -> {
      case get_codepoint(first) {
        c if c >= upper_a_charcode && c <= upper_z_charcode ->
          Some(upper_z_charcode - c + 2)
        c if c >= lower_a_charcode && c <= lower_z_charcode ->
          Some(c - lower_a_charcode + 2)
        _ -> None
      }
    }
  }
}

// check that the variable length integer encoding is a valid format
fn validate_integer_ok(i: String) -> Option(String) {
  let expected_length =
    i
    |> string.first
    |> result.unwrap("")
    |> string.to_option
    |> get_integer_length
    |> option.unwrap(-1)

  let ok = expected_length == i |> string.length
  ok |> bool.lazy_guard(return: fn() { Some(i) }, otherwise: fn() { None })
}

// -----------------------------------------------------------------------------------
// FRACTIONAL ENCODING
// -----------------------------------------------------------------------------------

// Returns a fractional key by finding the midpoint between a and b. Conducts appropriate checks
// for the order between a,b and the encoding of the values.
@internal
pub fn midpoint(
  a: Option(String),
  b: Option(String),
  encoding: String,
) -> String {
  let a = a |> option.unwrap("0")
  let b = b |> option.unwrap("")
  let #(prefix, a, b) = remove_common_prefix(a, b)
  case prefix {
    None -> {
      get_midpoint_between(a, b, encoding)
    }
    Some(prefix) ->
      prefix <> midpoint(a |> string.to_option, b |> string.to_option, encoding)
  }
}

// removes the common prefix between a, b if it exists. Returns #(Some(prefix), remainder of a, remainder of b) if prefix
// exists, or #(None, a, b)
fn remove_common_prefix(
  a: String,
  b: String,
) -> #(Option(String), String, String) {
  let prefix =
    list.zip(
      string.pad_end(a, to: b |> string.length, with: "0") |> string.split(""),
      b |> string.split(""),
    )
    |> list.fold_until("", fn(acc, p) {
      case p.0, p.1 {
        a, b if a == b -> list.Continue(acc <> a)
        _, _ -> list.Stop(acc)
      }
    })
  case prefix {
    "" -> #(None, a, b)
    prefix -> {
      let len = prefix |> string.length
      let a_rem = string.drop_start(a, up_to: len)
      let b_rem = string.drop_start(b, up_to: len)
      #(Some(prefix), a_rem, b_rem)
    }
  }
}

// calculates a midpoint between a and b.  a must be a non-empty string.  b may be the empty string. Requires a < b as checked
// by guard before calling this function
fn get_midpoint_between(a: String, b: String, encoding: String) -> String {
  let lut =
    list.zip(
      encoding |> string.split(""),
      list.range(0, encoding |> string.length),
    )
    |> dict.from_list

  // a must be a non-empty string
  let assert Ok(a_first) = a |> string.first
  let assert Ok(index_a) = lut |> dict.get(a_first)

  // b may be a null value.  If null, index is defined as one more than the
  // encoding length
  let b_first = b |> string.first |> result.unwrap("")
  let index_b =
    lut |> dict.get(b_first) |> result.unwrap(encoding |> string.length)

  // choose the midpoint between a, b indices
  let assert Ok(index_midpt) = index_a |> int.add(index_b) |> int.divide(2)

  case index_a, index_b, index_midpt {
    // if midpoint is greater than index of a, it means that a and b are not consecutive
    // so take the midpoint value
    index_a, _, index_midpt if index_midpt > index_a ->
      string.slice(encoding, at_index: index_midpt, length: 1)

    // if here, the values are consecutive, the right strategy depends on the length of b
    _, _, _ -> {
      case b |> string.length {
        // if b is too short, drop one character from a and calculate a new midpoint
        0 | 1 -> {
          let a_rem = a |> string.drop_start(up_to: 1) |> string.to_option
          a_first <> midpoint(a_rem, None, encoding)
        }
        // if b is multicharacters, then the first digit of b is between a and b
        _ -> b_first
      }
    }
  }
}

// -------------------------------------------------------------------------------
// Error Checks
// -------------------------------------------------------------------------------

fn order_ok(a: String, b: String) -> Result(Nil, KeyError) {
  let ok = string.compare(a, b) == order.Lt
  ok
  |> bool.lazy_guard(return: fn() { Ok(Nil) }, otherwise: fn() {
    Error(ErrWrongOrder)
  })
}

// all characters in the key must come from the encoding alphabet
fn encoding_ok(s: String, encoding: String) -> Result(Nil, KeyError) {
  s
  |> string.split("")
  |> list.map(fn(char) { string.contains(encoding, char) })
  |> list.fold_until(Ok(Nil), fn(_, res) {
    case res {
      False -> list.Stop(Error(ErrInvalid))
      True -> list.Continue(Ok(Nil))
    }
  })
}

// perform checks on the key to make sure that it is a valid encoding format
fn order_key_ok(key: String) -> Result(Nil, KeyError) {
  let int_part = get_integer_part(key)
  let fractional_part = case int_part {
    Some(int_key) -> string.drop_start(key, up_to: int_key |> string.length)
    // if integer part is bad, set it to 0 sentinel to fail next checks
    None -> "0"
  }

  let is_smallest_integer = key == smallest_integer
  let fractional_ends_in_zero = case fractional_part {
    "" -> False
    part -> part |> string.last |> result.unwrap("0") == "0"
  }

  let all_ok = !is_smallest_integer && !fractional_ends_in_zero

  all_ok
  |> bool.lazy_guard(return: fn() { Ok(Nil) }, otherwise: fn() {
    Error(ErrInvalid)
  })
}