///|
fn require_arity(
  name : String,
  actual : Int,
  expected : Int,
) -> Unit raise JmesPathError {
  if actual != expected {
    fail(
      "invalid_arity",
      0,
      "\{name} expects \{expected} argument(s), got \{actual}",
    )
  }
}

///|
fn require_min_arity(
  name : String,
  actual : Int,
  minimum : Int,
) -> Unit raise JmesPathError {
  if actual < minimum {
    fail(
      "invalid_arity",
      0,
      "\{name} expects at least \{minimum} argument(s), got \{actual}",
    )
  }
}

///|
fn[T] invalid_type(name : String) -> T raise JmesPathError {
  fail("invalid_type", 0, "invalid argument type for " + name)
}

///|
fn evaluate_args(
  args : Array[Expr],
  data : Json,
  state : EvalState,
) -> Array[Json] raise JmesPathError {
  let values : Array[Json] = []
  for arg in args {
    values.push(evaluate(arg, data, state))
  }
  values
}

///|
fn homogeneous_order(
  values : Array[Json],
  name : String,
  state : EvalState,
) -> Int raise JmesPathError {
  if values.is_empty() {
    return 0
  }
  match values[0] {
    Number(_, ..) => {
      for value in values {
        state.tick()
        if !(value is Number(_, ..)) {
          return invalid_type(name)
        }
      }
      1
    }
    String(_) => {
      for value in values {
        state.tick()
        if !(value is String(_)) {
          return invalid_type(name)
        }
      }
      2
    }
    _ => invalid_type(name)
  }
}

///|
fn compare_ordered(
  left : Json,
  right : Json,
  kind : Int,
  state : EvalState,
) -> Int raise JmesPathError {
  state.tick()
  if kind == 1 {
    guard left is Number(a, ..) && right is Number(b, ..) else { return 0 }
    if a < b {
      -1
    } else if a > b {
      1
    } else {
      0
    }
  } else {
    guard left is String(a) && right is String(b) else { return 0 }
    // JMESPath ordering is by Unicode code points, not UTF-16 code units.
    // Iterators avoid allocating a character array for every sort comparison.
    let left_chars = a.iter()
    let right_chars = b.iter()
    while true {
      state.tick()
      match (left_chars.next(), right_chars.next()) {
        (None, None) => return 0
        (None, Some(_)) => return -1
        (Some(_), None) => return 1
        (Some(x), Some(y)) => {
          if x.to_int() < y.to_int() {
            return -1
          }
          if x.to_int() > y.to_int() {
            return 1
          }
        }
      }
    } nobreak {
      0
    }
  }
}

///|
fn function_length(value : Json, state : EvalState) -> Json raise JmesPathError {
  match value {
    String(text) => {
      let mut count = 0
      for _ in text.iter() {
        state.tick()
        count += 1
      }
      Json::number(count.to_double())
    }
    Array(items) => Json::number(items.length().to_double())
    Object(entries) => Json::number(entries.length().to_double())
    _ => invalid_type("length")
  }
}

///|
fn function_contains(
  subject : Json,
  search : Json,
  state : EvalState,
) -> Json raise JmesPathError {
  match subject {
    String(text) =>
      match search {
        String(needle) => Json::boolean(text.contains(needle))
        _ => invalid_type("contains")
      }
    Array(items) => {
      for item in items {
        state.tick()
        if json_equal(item, search, state) {
          return Json::boolean(true)
        }
      }
      Json::boolean(false)
    }
    _ => invalid_type("contains")
  }
}

///|
fn function_numeric_array(
  name : String,
  value : Json,
  average : Bool,
  state : EvalState,
) -> Json raise JmesPathError {
  guard value is Array(items) else { return invalid_type(name) }
  let mut total = 0.0
  for item in items {
    state.tick()
    guard item is Number(number, ..) else { return invalid_type(name) }
    total += number
  }
  if average {
    if items.is_empty() {
      Json::null()
    } else {
      Json::number(total / items.length().to_double())
    }
  } else {
    Json::number(total)
  }
}

///|
fn function_extreme(
  name : String,
  value : Json,
  maximum : Bool,
  state : EvalState,
) -> Json raise JmesPathError {
  guard value is Array(items) else { return invalid_type(name) }
  if items.is_empty() {
    return Json::null()
  }
  let kind = homogeneous_order(items, name, state)
  let mut best = items[0]
  for i = 1; i < items.length(); i = i + 1 {
    let order = compare_ordered(items[i], best, kind, state)
    if (maximum && order > 0) || (!maximum && order < 0) {
      best = items[i]
    }
  }
  best
}

///|
fn function_sort(value : Json, state : EvalState) -> Json raise JmesPathError {
  guard value is Array(items) else { return invalid_type("sort") }
  state.add_results(items.length())
  let kind = homogeneous_order(items, "sort", state)
  Json::array(
    bounded_sort(items, (a, b) => compare_ordered(a, b, kind, state), state),
  )
}

///|
fn json_type(value : Json) -> String {
  match value {
    Null => "null"
    True | False => "boolean"
    Number(_, ..) => "number"
    String(_) => "string"
    Array(_) => "array"
    Object(_) => "object"
  }
}

///|
fn evaluate_expref_function(
  name : String,
  args : Array[Expr],
  data : Json,
  state : EvalState,
) -> Json raise JmesPathError {
  require_arity(name, args.length(), 2)
  let (array_expr, reference) = if name == "map" {
    guard args[0] is ExpressionRef(reference) else { return invalid_type(name) }
    (args[1], reference)
  } else {
    guard args[1] is ExpressionRef(reference) else { return invalid_type(name) }
    (args[0], reference)
  }
  guard evaluate(array_expr, data, state) is Array(items) else {
    return invalid_type(name)
  }
  if name == "map" {
    let mapped : Array[Json] = []
    for item in items {
      let value = evaluate(reference, item, state)
      state.add_results(1)
      mapped.push(value)
    }
    return Json::array(mapped)
  }
  if items.is_empty() {
    return if name == "sort_by" { Json::array([]) } else { Json::null() }
  }
  if name == "sort_by" {
    state.add_results(items.length())
  }
  let pairs : Array[(Json, Json)] = []
  for item in items {
    pairs.push((item, evaluate(reference, item, state)))
  }
  let keys = pairs.map(pair => pair.1)
  let kind = homogeneous_order(keys, name, state)
  if name == "sort_by" {
    let sorted = bounded_sort(
      pairs,
      (a, b) => compare_ordered(a.1, b.1, kind, state),
      state,
    )
    return Json::array(sorted.map(pair => pair.0))
  }
  let maximum = name == "max_by"
  let mut best = pairs[0]
  for i = 1; i < pairs.length(); i = i + 1 {
    let order = compare_ordered(pairs[i].1, best.1, kind, state)
    if (maximum && order > 0) || (!maximum && order < 0) {
      best = pairs[i]
    }
  }
  best.0
}

///|
fn evaluate_function(
  name : String,
  args : Array[Expr],
  data : Json,
  state : EvalState,
) -> Json raise JmesPathError {
  if name == "map" || name == "sort_by" || name == "max_by" || name == "min_by" {
    return evaluate_expref_function(name, args, data, state)
  }
  let values = evaluate_args(args, data, state)
  match name {
    "abs" => {
      require_arity(name, values.length(), 1)
      guard values[0] is Number(value, ..) else { return invalid_type(name) }
      Json::number(value.abs())
    }
    "avg" => {
      require_arity(name, values.length(), 1)
      function_numeric_array(name, values[0], true, state)
    }
    "ceil" => {
      require_arity(name, values.length(), 1)
      guard values[0] is Number(value, ..) else { return invalid_type(name) }
      Json::number(value.ceil())
    }
    "contains" => {
      require_arity(name, values.length(), 2)
      function_contains(values[0], values[1], state)
    }
    "ends_with" => {
      require_arity(name, values.length(), 2)
      guard values[0] is String(subject) && values[1] is String(suffix) else {
        return invalid_type(name)
      }
      Json::boolean(subject.has_suffix(suffix))
    }
    "floor" => {
      require_arity(name, values.length(), 1)
      guard values[0] is Number(value, ..) else { return invalid_type(name) }
      Json::number(value.floor())
    }
    "join" => {
      require_arity(name, values.length(), 2)
      guard values[0] is String(separator) && values[1] is Array(items) else {
        return invalid_type(name)
      }
      let strings : Array[String] = []
      let mut units = 0
      for item in items {
        state.tick()
        guard item is String(text) else { return invalid_type(name) }
        let separator_size = if strings.is_empty() {
          0
        } else {
          separator.length()
        }
        if separator_size > 16777216 - units {
          fail("limit_exceeded", 0, "join exceeds string size budget")
        }
        units += separator_size
        if text.length() > 16777216 - units {
          fail("limit_exceeded", 0, "join exceeds string size budget")
        }
        units += text.length()
        strings.push(text)
      }
      Json::string(strings.join(separator))
    }
    "keys" => {
      require_arity(name, values.length(), 1)
      guard values[0] is Object(entries) else { return invalid_type(name) }
      state.add_results(entries.length())
      let keys : Array[Json] = []
      for entry in entries {
        keys.push(Json::string(entry.0))
      }
      Json::array(keys)
    }
    "length" => {
      require_arity(name, values.length(), 1)
      function_length(values[0], state)
    }
    "max" => {
      require_arity(name, values.length(), 1)
      function_extreme(name, values[0], true, state)
    }
    "merge" => {
      require_min_arity(name, values.length(), 1)
      let result : Map[String, Json] = Map([])
      for value in values {
        guard value is Object(entries) else { return invalid_type(name) }
        for entry in entries {
          state.tick()
          if !result.contains(entry.0) {
            state.add_results(1)
          }
          result[entry.0] = entry.1
        }
      }
      Json::object(result)
    }
    "min" => {
      require_arity(name, values.length(), 1)
      function_extreme(name, values[0], false, state)
    }
    "not_null" => {
      require_min_arity(name, values.length(), 1)
      for value in values {
        if !(value is Null) {
          return value
        }
      }
      Json::null()
    }
    "reverse" => {
      require_arity(name, values.length(), 1)
      match values[0] {
        Array(items) => {
          state.add_results(items.length())
          Json::array(items.rev())
        }
        String(text) => {
          // Reject unaffordable copies before allocating output. UTF-16 units
          // conservatively cover every code point, including surrogate pairs.
          state.charge(text.length())
          let reversed = StringBuilder()
          for c in text.rev_iter() {
            reversed.write_char(c)
          }
          Json::string(reversed.to_string())
        }
        _ => invalid_type(name)
      }
    }
    "sort" => {
      require_arity(name, values.length(), 1)
      function_sort(values[0], state)
    }
    "starts_with" => {
      require_arity(name, values.length(), 2)
      guard values[0] is String(subject) && values[1] is String(prefix) else {
        return invalid_type(name)
      }
      Json::boolean(subject.has_prefix(prefix))
    }
    "sum" => {
      require_arity(name, values.length(), 1)
      function_numeric_array(name, values[0], false, state)
    }
    "to_array" => {
      require_arity(name, values.length(), 1)
      match values[0] {
        Array(_) => values[0]
        value => {
          state.add_results(1)
          Json::array([value])
        }
      }
    }
    "to_number" => {
      require_arity(name, values.length(), 1)
      match values[0] {
        Number(_, ..) => values[0]
        String(text) => {
          // Only a numeric scalar may reach the JSON parser. Parsing arbitrary
          // encoded arrays/objects here can allocate millions of discarded nodes.
          // Keep the parser responsible for exact JSON-number grammar.
          for c in text.iter() {
            state.tick()
            if !(is_digit(c) ||
              is_space(c) ||
              c == '-' ||
              c == '+' ||
              c == '.' ||
              c == 'e' ||
              c == 'E') {
              return Json::null()
            }
          }
          let parsed = @json.parse(text) catch { _ => return Json::null() }
          if parsed is Number(_, ..) {
            parsed
          } else {
            Json::null()
          }
        }
        _ => Json::null()
      }
    }
    "to_string" => {
      require_arity(name, values.length(), 1)
      match values[0] {
        String(_) => values[0]
        value => Json::string(stringify_bounded(value))
      }
    }
    "type" => {
      require_arity(name, values.length(), 1)
      Json::string(json_type(values[0]))
    }
    "values" => {
      require_arity(name, values.length(), 1)
      guard values[0] is Object(entries) else { return invalid_type(name) }
      state.add_results(entries.length())
      let result : Array[Json] = []
      for entry in entries {
        result.push(entry.1)
      }
      Json::array(result)
    }
    _ => fail("unknown_function", 0, "unknown function: " + name)
  }
}