string_width

import gleam/int
import gleam/list
import gleam/regex.{type Regex}
import gleam/string
import string_width/internal

// sources:
// UAX #11 east asian width: https://www.unicode.org/reports/tr11/
// Grapheme Clusters in Terminals: https://mitchellh.com/writing/grapheme-clusters-in-terminals
// get-east-asian-width: https://github.com/sindresorhus/get-east-asian-width
// string-width: https://github.com/sindresorhus/string-width
// ansi-regex: https://github.com/chalk/ansi-regex
// fast-string-truncated-width: https://github.com/fabiospampinato/fast-string-truncated-width/
// musl-libc wcwidth: https://git.musl-libc.org/cgit/musl/tree/src/ctype/wcwidth.c
// Mode 2027: https://github.com/contour-terminal/terminal-unicode-core
// https://gitlab.freedesktop.org/terminal-wg/specifications/-/issues/36

// v3:
// - we can optimise the top-level non-foldy things by stripping out things of known length beforehand
//   see the fstw regex magic
// - options builder, _with api
// - handle tabs
// - implement line wrapping? truncation? pad/alignment?
//   how  much would this break optimisations?
//   I should provide ways to implement layout, not do layout myself

/// Options to change the default behaviour of the functions in this library.
/// If you are unsure what to do here, passing an empty array is almost always the right call!
pub type Option {
  /// Most terminal emulators do not handle grapheme clusters well and will
  /// instead show their decomposition. To make sure a given string always fits
  /// even on those terminals, the functions in this package will copy this
  /// behaviour as well.
  ///
  /// If you pass this option, the returned numbers will more accurately represent
  /// the width of a string on a website or in an editor.
  ///
  /// **NOTE:** Grapheme support in terminals is highly experimental and subject
  /// to ongoing discussion! When working on CLIs or TUIs, you do not want to use
  /// this option most of the time.
  ///
  /// See also [Grapheme Clusters and Terminal Emulators](https://mitchellh.com/writing/grapheme-clusters-in-terminals)
  /// for a better explanation on how terminals behave.
  HandleGraphemeClusters
  /// Do not ignore ansi escape sequences, and count them as regular characters.
  ///
  /// You can pass this option as an optimisation if you are sure that your
  /// string doesn't contain any ansi escape codes.
  CountAnsiEscapeCodes
  /// Some characters are marked by Unicode as "ambiguous", meaning they may
  /// occupy 1 or 2 cells, depending on the context, current language, selected
  /// font, surrounding text, and more.
  ///
  /// Unicode recommends treating these characters as narrow by default,
  /// but you can change this behaviour using this option.
  AmbiguousAsWide
}

type Options {
  Options(
    count_ansi_escape_codes: Bool,
    ambiguous_as_wide: Bool,
    handle_grapheme_clusters: Bool,
  )
}

fn parse_options(options: List(Option)) -> Options {
  let defaults = Options(False, False, False)
  use record, option <- list.fold(options, defaults)
  case option {
    CountAnsiEscapeCodes -> Options(..record, count_ansi_escape_codes: True)
    AmbiguousAsWide -> Options(..record, ambiguous_as_wide: True)
    HandleGraphemeClusters -> Options(..record, handle_grapheme_clusters: True)
  }
}

/// Get the number of columns required to print a line in a terminal.
///
/// Line breaks are ignored. If the given string contains newlines,
/// the result is undefined.
///
/// ### Examples
///
/// ```gleam
/// line("äöüè", [])
/// // --> 4
///
/// line("안녕하세요", [])
/// // --> 10
///
/// line("👩‍👩‍👦‍👦", [])
/// // --> 8
///
/// line("👩‍👩‍👦‍👦", [HandleGraphemeClusters])
/// // --> 2
///
/// line("\u{1B}[31mhello\u{1B}[39m", [])
/// // --> 5
/// ```
pub fn line(str: String, options: List(Option)) -> Int {
  let options = parse_options(options)
  // we do not need to fold over ansi codes if we gonna ignore them anyways
  let str = case options.count_ansi_escape_codes {
    True -> str
    False -> regex.replace(each: ansi_re(), in: str, with: "")
  }

  // strip out ASCII, those cannot form a grapheme cluster and are always narrow.
  // this does helps mostly in grapheme cluster mode.
  let encoded_length_before = encoded_byte_size(str)
  let str = regex.replace(each: ascii_re(), in: str, with: "")
  let ascii_length = encoded_length_before - encoded_byte_size(str)

  // we do not need to pas count_ansi_escape_codes forward, we have already handled them.
  do_line(options, str) + ascii_length
}

fn do_line(options: Options, str: String) -> Int {
  use sum, _chr, width <- do_fold(options, str, 0, [], 0)
  sum + width
}

/// The required number of rows and columns to display the given text.
///
/// The number of rows is equal to the number of lines, while the number of
/// columns represents the maximum line width.
///
/// ### Examples
///
/// ```gleam
/// dimensions("안녕하세요", [])
/// // --> #(1, 10)
///
/// dimensions("hello,\n안녕하세요", [])
/// // --> #(2, 10)
/// ```
pub fn dimensions(str: String, options: List(Option)) -> #(Int, Int) {
  let #(rows, cols, cols_max) = {
    use #(rows, cols, cols_max), chr, width <- fold(str, options, #(1, 0, 0))
    case chr {
      "\n" -> #(rows + 1, 0, int.max(cols, cols_max))
      _ -> #(rows, cols + width, cols_max)
    }
  }

  #(rows, int.max(cols, cols_max))
}

/// Iterate over the measured components of a string. Components are either
/// graphemes or codepoints, depending on the `HandleGraphemeClusters` option,
/// or other undivisible sequences, like ANSI escape codes.
///
/// This is a lower-level utility compared to the others in this package. It does
/// not by itself keep track of any additional state, and does not for example
/// handle newlines or tabs. Instead, you can use fold to implement all kinds of
/// higher-level layout primitives.
///
/// Concatenating all components is guaranteed to produce the original string.
///
/// ## Examples
///
/// ```gleam
/// // A slower string_width.line implementation that doesn't handle tabs
/// fold("hello", [], from: 0, with: fn(so_far, _chr, width) { width + so_far })
/// // --> 5
/// ```
///
/// ```gleam
/// // truncate a string after 50 characters, but keep all ansi sequences.
/// use #(total, acc), chr, width <- fold(input, [], from: #(0, ""))
/// case total >= 50 {
///   True -> case chr {
///     "\u{1b}" <> _ -> #(total, acc <> chr)
///      _ -> #(total, acc)
///   }
///
///   False -> #(total + width, acc <> chr)
/// }
/// ```
pub fn fold(
  over string: String,
  using options: List(Option),
  from state: state,
  with fun: fn(state, String, Int) -> state,
) -> state {
  let options = parse_options(options)
  let ansi_ranges = case options.count_ansi_escape_codes {
    True -> []
    False -> scan_poslen(ansi_re(), string)
  }
  do_fold(options, string, 0, ansi_ranges, state, fun)
}

fn do_fold(
  options: Options,
  string: String,
  offset: Int,
  ansi_ranges: List(#(Int, Int)),
  state: state,
  fun: fn(state, String, Int) -> state,
) -> state {
  case ansi_ranges {
    [#(ansi_start, ansi_length), ..ansi_ranges] -> {
      let #(before, at_ansi) = unsafe_split(string, ansi_start - offset)
      let #(ansi_slice, string) = unsafe_split(at_ansi, ansi_length)

      let state =
        state
        |> do_fold_characters(options, before, _, fun)
        |> fun(ansi_slice, 0)

      let offset = ansi_start + ansi_length
      do_fold(options, string, offset, ansi_ranges, state, fun)
    }
    [] -> do_fold_characters(options, string, state, fun)
  }
}

fn do_fold_characters(
  options: Options,
  string: String,
  state: state,
  fun: fn(state, String, Int) -> state,
) -> state {
  case options.handle_grapheme_clusters {
    True -> do_fold_graphemes(options, string, state, fun)
    False -> do_fold_codepoints(options, string, state, fun)
  }
}

fn do_fold_codepoints(
  options: Options,
  string: String,
  state: state,
  fun: fn(state, String, Int) -> state,
) -> state {
  use state, cp <- fold_codepoints(string, state)
  let width = wcwidth(options, cp)
  fun(state, utf_codepoint_to_string(cp), width)
}

fn do_fold_graphemes(
  options: Options,
  string: String,
  state: state,
  fun: fn(state, String, Int) -> state,
) -> state {
  use state, grapheme <- fold_graphemes(string, state)
  let width = do_grapheme_cluster(options, grapheme)
  fun(state, grapheme, width)
}

/// Estimate the required width for a single grapheme cluster.
///
/// **Note:** Grapheme cluster handling in terminals is highly experimental and
/// subject to ongoing discussions, and very inconsistent across different
/// terminal emulators. The exact width is context-dependent, so this value may
/// may exactly match what you might expect!
pub fn grapheme_cluster(grapheme: String, options: List(Option)) -> Int {
  do_grapheme_cluster(parse_options(options), grapheme)
}

fn do_grapheme_cluster(options: Options, string: String) -> Int {
  // The Unicode Core spec (https://github.com/contour-terminal/terminal-unicode-core)
  // proposing Mode 2027 says the following:
  //
  // > Therefore, extending a grapheme cluster with consecutively added codepoints
  // > will not move the cursor except for variation selector 16 (VS16) that may
  // > have caused the width of the grapheme cluster to change to wide (2 grid cells).
  //
  // I do not agree.
  //
  // Many grapheme clusters include multiple letters with a width, which cannot
  // be printed and combined into a single cell. (examples: च्छे षि ‍র্য ﾊﾞ)
  //
  // Therefore, in addition to the rules in this proposal, having 2 or more
  // non-modifier codepoints with width in a cluster forces it to wide.
  //
  // This more accurately represents what font shaping would produce in
  // non-terminal applications, and not a single terminal in my tests implements
  // the aformentioned spec correctly (in my understanding) anyways.
  use max, cp <- fold_codepoints(string, 0)
  case cp {
    // unicode core: VS16 forces a cluster to be wide
    0xfe0f -> 2
    _ -> {
      let width = wcwidth(options, cp)
      case width > 0 && max > 0 {
        // extension: 2 or more codepoints with width force a cluster to be wide
        True ->
          case is_spacing_mark(cp) {
            True -> int.max(max, width)
            False -> 2
          }
        False -> int.max(max, width)
      }
    }
  }
}

/// Estimate the required width of a single unicode code point.
///
/// If you encounter a mismatch that is consistent across multiple terminal
/// emulators, please open an issue or ping me on Discord!
pub fn codepoint(chr: UtfCodepoint, options: List(Option)) -> Int {
  wcwidth(parse_options(options), string.utf_codepoint_to_int(chr))
}

///
pub fn int_codepoint(cp: Int, options: List(Option)) -> Int {
  wcwidth(parse_options(options), cp)
}

fn wcwidth(options: Options, cp: Int) -> Int {
  // see: https://git.musl-libc.org/cgit/musl/tree/src/ctype/wcwidth.c
  case cp {
    // cheap optimization for ascii
    _ if cp <= 0x1f -> 0
    _ if cp <= 0x7e -> 1
    _ ->
      case is_ignorable(cp) {
        True -> 0
        False ->
          case is_ambiguous(cp) {
            True ->
              case options.ambiguous_as_wide {
                True -> 2
                False -> 1
              }
            False ->
              case is_wide(cp) {
                True -> 2
                False -> 1
              }
          }
      }
  }
}

// -- TABLE ACCESSORS ----------------------------------------------------------

// the data for these functions is generated by the generate.gleam script.
// the general idea is heavily inspired by the tables used in musl libc.

fn is_ignorable(cp: Int) -> Bool {
  case cp <= 0x1ffff {
    True -> table_lookup(internal.ignorable, cp)
    False -> cp >= 0xe0000 && cp <= 0xe0fff
  }
}

fn is_ambiguous(cp: Int) -> Bool {
  case cp <= 0xffff {
    True -> table_lookup(internal.ambiguous, cp)
    False -> {
      { cp >= 0x1f100 && cp <= 0x1f10a }
      || { cp >= 0x1f110 && cp <= 0x1f12d }
      || { cp >= 0x1f130 && cp <= 0x1f169 }
      || { cp >= 0x1f170 && cp <= 0x1f18d }
      || { cp >= 0x1f18f && cp <= 0x1f190 }
      || { cp >= 0x1f19b && cp <= 0x1f1ac }
      || { cp >= 0xe0100 && cp <= 0xe01ef }
      || { cp >= 0xf0000 && cp <= 0xffffd }
      || { cp >= 0x100000 && cp <= 0x10FFFD }
    }
  }
}

fn is_wide(cp: Int) -> Bool {
  case cp <= 0x1ffff {
    True -> table_lookup(internal.wide, cp)
    False -> {
      // NOTE: this assumes we checked ambiguous first!
      cp >= 0x20000 && cp <= 0x40000
    }
  }
}

fn is_spacing_mark(cp: Int) -> Bool {
  case cp <= 0x1ffff {
    True -> table_lookup(internal.spacing_mark, cp)
    False -> False
  }
}

// -- FFI ----------------------------------------------------------------------

@external(javascript, "./string_width_ffi.mjs", "table_lookup")
fn table_lookup(table: BitArray, value: Int) -> Bool {
  let hi = int.bitwise_shift_right(value, 8)
  let md = int.bitwise_shift_right(int.bitwise_and(value, 0xff), 3)
  let lo = int.bitwise_and(value, 7)

  let lvl1_skip = hi * 8
  let assert <<_:size(lvl1_skip), lvl1:int, _:bits>> = table

  // + 7 - lo, then selecting a single bit saves us a bitshift and a modulo
  // this is equivalent to ((lvl3 >> lo) & 1)
  // unfortunately not supported on javascript, but we FFI anyways to go fast.
  let lvl2_skip = { lvl1 * 32 + md } * 8 + 7 - lo
  let assert <<_:size(lvl2_skip), lvl3:size(1), _:bits>> = table
  lvl3 != 0
}

@external(javascript, "./string_width_ffi.mjs", "fold_graphemes")
fn fold_graphemes(
  str: String,
  state: state,
  fun: fn(state, String) -> state,
) -> state {
  case string.pop_grapheme(str) {
    Ok(#(grapheme, str)) -> fold_graphemes(str, fun(state, grapheme), fun)
    Error(Nil) -> state
  }
}

@external(erlang, "string_width_ffi", "fold_codepoints")
@external(javascript, "./string_width_ffi.mjs", "fold_codepoints")
fn fold_codepoints(
  str: String,
  state: state,
  fun: fn(state, Int) -> state,
) -> state {
  use state, cp <- list.fold(string.to_utf_codepoints(str), state)
  fun(state, string.utf_codepoint_to_int(cp))
}

@external(erlang, "string_width_ffi", "ansi_re")
@external(javascript, "./string_width_ffi.mjs", "ansi_re")
fn ansi_re() -> Regex

@external(erlang, "string_width_ffi", "ascii_re")
@external(javascript, "./string_width_ffi.mjs", "ascii_re")
fn ascii_re() -> Regex

@external(erlang, "string_width_ffi", "scan_poslen")
@external(javascript, "./string_width_ffi.mjs", "scan_poslen")
fn scan_poslen(regex: Regex, str: String) -> List(#(Int, Int))

@external(erlang, "erlang", "byte_size")
@external(javascript, "./string_width_ffi.mjs", "encoded_byte_size")
fn encoded_byte_size(str: String) -> int

@external(erlang, "erlang", "split_binary")
@external(javascript, "./string_width_ffi.mjs", "unsafe_split")
fn unsafe_split(str: String, at: Int) -> #(String, String)

@external(erlang, "string_width_ffi", "utf_codepoint_to_string")
@external(javascript, "./string_width_ffi.mjs", "utf_codepoint_to_string")
fn utf_codepoint_to_string(cp: Int) -> String