///|
priv enum Argument {
  Value(Json)
  Reference(Node)
}

///|
fn signature(name : String) -> Array[String]? {
  match name {
    "abs" | "ceil" | "floor" => Some(["number"])
    "avg" | "sum" => Some(["array[number]"])
    "contains" => Some(["string|array", "any"])
    "starts_with" | "ends_with" => Some(["string", "string"])
    "join" => Some(["string", "array[string]"])
    "keys" | "values" | "merge" => Some(["object"])
    "length" => Some(["string|array|object"])
    "map" => Some(["expref", "array"])
    "max" | "min" | "sort" => Some(["array[sortable]"])
    "max_by" | "min_by" | "sort_by" => Some(["array", "expref"])
    "reverse" => Some(["string|array"])
    "not_null" | "to_array" | "to_number" | "to_string" | "type" =>
      Some(["any"])
    _ => None
  }
}

///|
fn accepts(value : Json, expected : String) -> Bool {
  if expected == "any" {
    return true
  }
  if expected.contains("|") {
    return expected.split("|").any(t => accepts(value, t.to_owned()))
  }
  match expected {
    "array[number]" =>
      match value {
        Array(xs) => xs.iter().all(x => x is Number(_, ..))
        _ => false
      }
    "array[string]" =>
      match value {
        Array(xs) => xs.iter().all(x => x is String(_))
        _ => false
      }
    "array[sortable]" =>
      match value {
        Array(xs) =>
          xs.is_empty() ||
          xs.iter().all(x => x is Number(_, ..)) ||
          xs.iter().all(x => x is String(_))
        _ => false
      }
    _ => json_type(value) == expected
  }
}

///|
fn[T] type_error(
  name : String,
  expected : String,
  span : Span,
) -> T raise JmespathError {
  raise JmespathError({
    code: InvalidType,
    span,
    message: "I expected \{expected} in JMESPath function \{name}.",
  })
}

///|
fn json_argument(
  arg : Argument,
  name : String,
  span : Span,
) -> Json raise JmespathError {
  match arg {
    Value(value) => value
    Reference(_) => type_error(name, "a JSON argument", span)
  }
}

///|
fn finite_number(value : Double, span : Span) -> Json raise JmespathError {
  if value.is_nan() || value.is_inf() {
    problem(
      InvalidValue,
      span,
      "I expected a finite JSON number from the JMESPath function.",
    )
  }
  Json::number(value)
}

///|
/// Compare numbers as the Smithy reference does (see `number_order`), and
/// strings by UTF-16 code unit, as Java's `String.compareTo` does.
fn compare_json(a : Json, b : Json) -> Int32 {
  match (a, b) {
    (Number(x, repr=x_repr), Number(y, repr=y_repr)) =>
      number_order(x, x_repr, y, y_repr)
    (String(a), String(b)) => {
      let n = a.length().min(b.length())
      for i in 0.. 0
  }
}

///|
fn check_arity(
  name : String,
  count : Int32,
  span : Span,
) -> Array[String] raise JmespathError {
  guard signature(name) is Some(types) else {
    problem(
      UnknownFunction,
      span,
      "I could not resolve JMESPath function \{name}.",
    )
    return []
  }
  let variadic = name == "merge" || name == "not_null"
  if count < types.length() || (!variadic && count != types.length()) {
    problem(
      InvalidArity,
      span,
      "I expected \{types.length()} argument(s) for JMESPath function \{name}, but found \{count}.",
    )
  }
  types
}

///|
fn check_arguments(
  name : String,
  nodes : Array[Node],
  args : Array[Argument],
  span : Span,
) -> Unit raise JmespathError {
  let types = check_arity(name, args.length(), span)
  let variadic = name == "merge" || name == "not_null"
  for i, arg in args {
    let expected = types[if variadic { 0 } else { i }]
    let valid = match arg {
      Reference(_) => expected == "expref"
      Value(value) => accepts(value, expected)
    }
    if !valid {
      type_error(name, expected, nodes[i].span)
    }
  }
}

///|
fn call(
  name : String,
  nodes : Array[Node],
  args : Array[Argument],
  span : Span,
) -> Json raise JmespathError {
  // Typed arguments have been checked above. Keep JSON and expression arguments
  // separate so an expression reference can never leak into a JSON result.
  let first = match args[0] {
    Value(value) => value
    Reference(_) => Json::null()
  }
  match name {
    "type" => Json::string(json_type(first))
    "to_array" =>
      match first {
        Array(_) => first
        _ => Json::array([first])
      }
    "to_string" =>
      match first {
        String(_) => first
        _ => Json::string(first.stringify())
      }
    "to_number" =>
      match first {
        Number(n, ..) => {
          ignore(finite_number(n, span))
          // Keep the source token, so the number keeps its exact value.
          first
        }
        String(s) => to_number_of_string(s)
        _ => Json::null()
      }
    "not_null" => {
      for arg in args {
        let value = json_argument(arg, name, span)
        if value != Json::null() {
          return value
        }
      }
      Json::null()
    }
    "abs" | "ceil" | "floor" => {
      guard first is Number(n, ..) else {
        return type_error(name, "number", span)
      }
      finite_number(
        match name {
          "abs" => n.abs()
          "ceil" => n.ceil()
          _ => n.floor()
        },
        span,
      )
    }
    "length" => {
      let n = match first {
        String(s) => s.iter().count()
        Array(xs) => xs.length()
        Object(xs) => xs.length()
        _ => 0
      }
      Json::number(n.to_double())
    }
    "starts_with" | "ends_with" => {
      guard first is String(s) &&
        json_argument(args[1], name, span) is String(needle) else {
        return type_error(name, "string", span)
      }
      Json::boolean(
        if name == "starts_with" {
          s.has_prefix(needle)
        } else {
          s.has_suffix(needle)
        },
      )
    }
    "contains" => {
      let needle = json_argument(args[1], name, span)
      match first {
        String(s) =>
          match needle {
            String(needle) => Json::boolean(s.contains(needle))
            _ => type_error(name, "a string search value", nodes[1].span)
          }
        Array(xs) => Json::boolean(xs.iter().any(x => json_equal(x, needle)))
        _ => type_error(name, "string or array", span)
      }
    }
    "reverse" =>
      match first {
        String(s) =>
          Json::string(String::from_iter(s.iter().to_array().rev_iter()))
        Array(xs) => Json::array(xs.rev_iter().to_array())
        _ => type_error(name, "string or array", span)
      }
    "keys" | "values" => {
      guard first is Object(fields) else {
        return type_error(name, "object", span)
      }
      let out = []
      for key, value in fields {
        out.push(if name == "keys" { Json::string(key) } else { value })
      }
      Json::array(out)
    }
    "merge" => {
      let out : Map[String, Json] = Map([])
      for arg in args {
        guard json_argument(arg, name, span) is Object(fields) else {
          return type_error(name, "object", span)
        }
        for key, value in fields {
          out[key] = value
        }
      }
      Json::object(out)
    }
    "join" => {
      guard first is String(separator) &&
        json_argument(args[1], name, span) is Array(xs) else {
        return type_error(name, "string and array[string]", span)
      }
      let strings = []
      for value in xs {
        if value is String(s) {
          strings.push(s)
        }
      }
      Json::string(strings.join(separator))
    }
    "sum" | "avg" => {
      guard first is Array(xs) else {
        return type_error(name, "array[number]", span)
      }
      let mut sum = 0.0
      for x in xs {
        if x is Number(n, ..) {
          sum += n
        }
      }
      if name == "avg" && xs.is_empty() {
        Json::null()
      } else {
        finite_number(
          if name == "sum" {
            sum
          } else {
            sum / xs.length().to_double()
          },
          span,
        )
      }
    }
    "sort" | "min" | "max" => {
      guard first is Array(xs) else {
        return type_error(name, "an array of numbers or strings", span)
      }
      let xs = xs.copy()
      if name == "sort" {
        xs.sort_by(compare_json)
        Json::array(xs)
      } else {
        let mut best = Json::null()
        for x in xs {
          let order = compare_json(x, best)
          if best == Json::null() ||
            (name == "min" && order < 0) ||
            (name == "max" && order > 0) {
            best = x
          }
        }
        best
      }
    }
    _ => Json::null()
  }
}

///|
fn apply_keys(
  name : String,
  xs : Array[Json],
  keys : Array[Json],
  span : Span,
) -> Json raise JmespathError {
  let keyed : Array[(Json, Json, Int32)] = []
  let mut key_type = ""
  for i, key in keys {
    let kind = json_type(key)
    if (kind != "number" && kind != "string") ||
      (key_type != "" && key_type != kind) {
      return type_error(
        name, "uniform number or string expression results", span,
      )
    }
    key_type = kind
    keyed.push((key, xs[i], i))
  }
  if name == "sort_by" {
    keyed.sort_by(fn(a, b) {
      let order = compare_json(a.0, b.0)
      if order == 0 {
        a.2.compare(b.2)
      } else {
        order
      }
    })
    Json::array(keyed.map(entry => entry.1))
  } else {
    let mut best : (Json, Json, Int32)? = None
    for entry in keyed {
      match best {
        None => best = Some(entry)
        Some(old) => {
          let order = compare_json(entry.0, old.0)
          if (name == "min_by" && order < 0) || (name == "max_by" && order > 0) {
            best = Some(entry)
          }
        }
      }
    }
    best.map(entry => entry.1).unwrap_or(Json::null())
  }
}

///|
/// `to_number` of a string, as in the Smithy reference: a string with `.`,
/// `e` or `E` is a double, and any other string is a 64-bit integer or null.
fn to_number_of_string(s : String) -> Json {
  if is_decimal_token(s) {
    guard java_double(s) is Some(n) else { return Json::null() }
    // Core Json can hold non-finite doubles, as can arithmetic results here.
    // Keep overflow as a Double, without a decimal token, to match Smithy.
    if n.is_inf() {
      return Json::number(n)
    }
    // A decimal token that a double holds makes the number a double.
    let text = n.to_string()
    return Json::number(
      n,
      repr=if is_decimal_token(text) { text } else { text + ".0" },
    )
  }
  let digits = if s.has_prefix("-") || s.has_prefix("+") {
    s[1:].to_owned()
  } else {
    s
  }
  guard !digits.is_empty() && digits.iter().all(c => c >= '0' && c <= '9') else {
    return Json::null()
  }
  let token = if s.has_prefix("+") { digits } else { s }
  let whole = BigInt::from_string(token)
  guard whole >= long_min && whole <= long_max else { return Json::null() }
  let n = @string.parse_double(token) catch { _ => return Json::null() }
  Json::number(n, repr=token)
}