#include <rlang.h>
#include "globals.h"
#include "internal.h"
#include "nse-inject.h"
#include "utils.h"

#include "decl/arg-decl.h"


// Capture ----------------------------------------------------------------

static
r_obj* capture(r_obj* sym, r_obj* frame, r_obj** arg_env) {
  static r_obj* capture_call = NULL;
  if (!capture_call) {
    r_obj* args = KEEP(r_new_node(r_null, r_null));
    capture_call = r_new_call(rlang_ns_get("captureArgInfo"), args);
    r_preserve(capture_call);
    r_mark_shared(capture_call);
    FREE(1);
  }

  if (r_typeof(sym) != SYMSXP) {
    r_abort("`arg` must be a symbol");
  }

  r_node_poke_cadr(capture_call, sym);
  r_obj* arg_info = KEEP(r_eval(capture_call, frame));
  r_obj* expr = r_list_get(arg_info, 0);
  r_obj* env = r_list_get(arg_info, 1);

  // Unquoting rearranges the expression
  // FIXME: Only duplicate the call tree, not the leaves
  expr = KEEP(r_copy(expr));
  expr = call_interp(expr, env);

  if (arg_env) {
    *arg_env = env;
  }

  FREE(2);
  return expr;
}

r_obj* ffi_enexpr(r_obj* sym, r_obj* frame) {
  return capture(sym, frame, NULL);
}
r_obj* ffi_ensym(r_obj* sym, r_obj* frame) {
  r_obj* expr = capture(sym, frame, NULL);

  if (is_quosure(expr)) {
    expr = quo_get_expr(expr);
  }

  switch (r_typeof(expr)) {
  case R_TYPE_symbol:
    break;
  case R_TYPE_character:
    if (r_length(expr) == 1) {
      KEEP(expr);
      expr = r_sym(r_chr_get_c_string(expr, 0));
      FREE(1);
      break;
    }
    // else fallthrough
  default:
    // FIXME: Should call `abort_coercion()`
    r_abort("Can't convert to a symbol.");
  }

  return expr;
}


r_obj* ffi_enquo(r_obj* sym, r_obj* frame) {
  r_obj* env;
  r_obj* expr = KEEP(capture(sym, frame, &env));
  r_obj* quo = forward_quosure(expr, env);
  FREE(1);
  return quo;
}


// Match ------------------------------------------------------------------

static
int arg_match(r_obj* arg,
              r_obj* values,
              struct r_lazy error_arg,
              struct r_lazy error_call,
              struct r_lazy call) {
  if (r_typeof(values) != R_TYPE_character) {
    r_abort_lazy_call(call, "`values` must be a character vector.");
  }

  int values_len = r_length(values);
  if (values_len == 0) {
    r_abort_lazy_call(call, "`values` must have at least one element.");
  }

  switch (r_typeof(arg)) {
  case R_TYPE_character:
    break;
  case R_TYPE_string:
    return arg_match1(arg, values, error_arg, error_call);
  case R_TYPE_symbol:
    return arg_match1(r_sym_string(arg), values, error_arg, error_call);
  default:
    r_abort_lazy_call(error_call,
                      "%s must be a string or character vector.",
                      r_format_lazy_error_arg(error_arg));
  }

  int arg_len = r_length(arg);

  if (arg_len == 1) {
    return arg_match1(r_chr_get(arg, 0), values, error_arg, error_call);
  }

  if (arg_len != values_len) {
    r_abort_lazy_call(call, "`arg` must be a string or have the same length as `values`.");
  }

  r_obj* const* v_values = r_chr_cbegin(values);
  r_obj* const* v_arg = r_chr_cbegin(arg);

  // Same-length vector: must be identical, we allow changed order.
  int i = 0;
  for (; i < arg_len; ++i) {
    if (v_arg[i] != v_values[i]) {
      break;
    }
  }

  // Elements are identical, return first
  if (i == arg_len) {
    return 0;
  }

  r_obj* my_values = KEEP(r_clone(values));
  r_obj* const * v_my_values = r_chr_cbegin(my_values);

  // Invariant: my_values[i:(len-1)] contains the values we haven't matched yet
  for (; i < arg_len; ++i) {
    r_obj* current_arg = v_arg[i];
    if (current_arg == v_my_values[i]) {
      continue;
    }

    bool matched = false;
    for (int j = i + 1; j < arg_len; ++j) {
      if (current_arg == v_my_values[j]) {
        matched = true;

        // Replace matched value by the element that failed to match at this iteration
        r_chr_poke(my_values, j, v_my_values[i]);
        break;
      }
    }

    if (!matched) {
      r_eval_with_wxyz(stop_arg_match_call,
                       arg,
                       values,
                       KEEP(lazy_wrap_chr(error_arg)),
                       KEEP(r_lazy_eval(error_call)),
                       rlang_ns_env);
      r_stop_unreachable();
    }
  }

  r_obj* first_elt = r_chr_get(arg, 0);
  for (i = 0; i < values_len; ++i) {
    if (first_elt == v_values[i]) {
      FREE(1);
      return i;
    }
  }

  r_stop_unreachable();
}

int arg_match_legacy(r_obj* arg,
                     r_obj* values,
                     r_obj* error_arg,
                     r_obj* error_call) {
  struct r_lazy lazy_error_arg = { error_arg, r_null };
  struct r_lazy lazy_error_call = { error_call, r_null };

  return arg_match(arg,
                   values,
                   lazy_error_arg,
                   lazy_error_call,
                   r_lazy_null);
}

static
int arg_match1(r_obj* arg,
               r_obj* values,
               struct r_lazy error_arg,
               struct r_lazy error_call) {
  // Simple case: one argument, we check if it's one of the values
  r_obj* const* v_values = r_chr_cbegin(values);
  int n_values = r_length(values);

  for (int i = 0; i < n_values; ++i) {
    if (arg == v_values[i]) {
      return i;
    }
  }

  r_obj* ffi_error_call = r_lazy_eval(error_call);
  if (ffi_error_call == r_missing_arg) {
    // Replace `error_call` by environment on the stack because
    // `r_eval_with_` evaluates in an out-of-stack mask
    ffi_error_call = r_peek_frame();
  }
  KEEP(ffi_error_call);

  r_eval_with_wxyz(stop_arg_match_call,
                   KEEP(wrap_chr(arg)),
                   values,
                   KEEP(lazy_wrap_chr(error_arg)),
                   ffi_error_call,
                   rlang_ns_env);
  r_stop_unreachable();
}

static
r_obj* wrap_chr(r_obj* arg) {
  switch (arg_match_arg_nm_type(arg)) {
  case R_TYPE_string:
    return r_str_as_character(arg);
  case R_TYPE_symbol:
    return r_sym_as_utf8_character(arg);
  case R_TYPE_character:
    return arg;
  default:
    r_stop_unreachable();
  }
}

static
r_obj* lazy_wrap_chr(struct r_lazy arg) {
  r_obj* ffi_arg = KEEP(r_lazy_eval(arg));
  r_obj* out = wrap_chr(ffi_arg);
  FREE(1);
  return out;
}

static
enum r_type arg_match_arg_nm_type(r_obj* arg_nm) {
  switch (r_typeof(arg_nm)) {
  case R_TYPE_string: return R_TYPE_string;
  case R_TYPE_symbol: return R_TYPE_symbol;
  case R_TYPE_character:
    if (r_is_string(arg_nm)) {
      return R_TYPE_character;
    }
    // else fallthrough;
  default:
      r_abort("`arg_nm` must be a string or symbol.");
  }
}


int cci_arg_match(r_obj* arg,
                  r_obj* values,
                  struct r_lazy error_arg,
                  struct r_lazy error_call) {
  return arg_match(arg, values, error_arg, error_call, r_lazy_null);
}

r_obj* ffi_arg_match0(r_obj* args) {
  args = r_node_cdr(args);

  r_obj* arg = r_node_car(args); args = r_node_cdr(args);
  r_obj* values = r_node_car(args); args = r_node_cdr(args);
  r_obj* frame = r_node_car(args);

  struct r_lazy error_arg = { .x = syms.arg_nm, .env = frame };
  struct r_lazy error_call = { .x = r_syms.error_call, .env = frame };
  struct r_lazy call = { .x = frame, .env = r_null };

  int i = arg_match(arg, values, error_arg, error_call, call);
  return r_str_as_character(r_chr_get(values, i));
}


void rlang_init_arg(r_obj* ns) {
  stop_arg_match_call = r_parse("stop_arg_match(w, values = x, error_arg = y, error_call = z)");
  r_preserve(stop_arg_match_call);
}

static r_obj* stop_arg_match_call = NULL;
