#include <rlang.h>
#include "env.h"
#include "internal.h"
#include "quo.h"
#include "utils.h"

static r_obj* quo_mask_flag_sym = NULL;
static r_obj* data_mask_flag_sym = NULL;

enum rlang_mask_type {
    RLANG_MASK_DATA,    // Full data mask
    RLANG_MASK_QUOSURE, // Quosure mask with only `~` binding
    RLANG_MASK_NONE
};

struct rlang_mask_info {
    r_obj* mask;
    enum rlang_mask_type type;
};

static struct rlang_mask_info mask_info(r_obj* mask) {
    if (r_typeof(mask) != R_TYPE_environment) {
        return (struct rlang_mask_info) {r_null, RLANG_MASK_NONE};
    }

    r_obj* flag;
    r_obj* where;

    where = r_env_until(mask, data_mask_flag_sym, r_envs.empty);
    if (where != r_envs.empty) {
        flag = r_env_get(where, data_mask_flag_sym);
        return (struct rlang_mask_info) {flag, RLANG_MASK_DATA};
    }

    where = r_env_until(mask, quo_mask_flag_sym, r_envs.empty);
    if (where != r_envs.empty) {
        flag = r_env_get(where, quo_mask_flag_sym);
        return (struct rlang_mask_info) {flag, RLANG_MASK_QUOSURE};
    }

    return (struct rlang_mask_info) {r_null, RLANG_MASK_NONE};
}

static r_obj* data_pronoun_class = NULL;
static r_obj* ctxt_pronoun_class = NULL;
static r_obj* data_mask_env_sym = NULL;

static r_obj* rlang_new_data_pronoun(r_obj* mask) {
    r_obj* pronoun = KEEP(r_alloc_list(1));

    r_list_poke(pronoun, 0, mask);
    r_attrib_poke(pronoun, r_syms.class_, data_pronoun_class);

    FREE(1);
    return pronoun;
}
static r_obj* rlang_new_ctxt_pronoun(r_obj* top) {
    r_obj* pronoun = KEEP(r_alloc_empty_environment(r_env_parent(top)));

    r_attrib_poke(pronoun, r_syms.class_, ctxt_pronoun_class);

    FREE(1);
    return pronoun;
}

void poke_ctxt_env(r_obj* mask, r_obj* env) {
    if (!r_env_has(mask, data_mask_env_sym)) {
        r_abort("Internal error: Can't find context pronoun in data mask");
    }

    r_obj* ctxt_pronoun = KEEP(r_env_get(mask, data_mask_env_sym));
    r_env_poke_parent(ctxt_pronoun, env);
    FREE(1);
}

static r_obj* empty_names_chr;

static void check_unique_names(r_obj* x) {
    // Allow empty lists
    if (!r_length(x)) {
        return;
    }

    r_obj* names = KEEP(r_names(x));
    if (names == r_null) {
        r_abort("`data` must be uniquely named but does not have names");
    }
    if (vec_find_first_duplicate(names, empty_names_chr, NULL)) {
        r_abort("`data` must be uniquely named but has duplicate columns");
    }
    FREE(1);
}
r_obj* ffi_as_data_pronoun(r_obj* x) {
    int n_kept = 0;

    switch (r_typeof(x)) {
    case R_TYPE_logical:
    case R_TYPE_integer:
    case R_TYPE_double:
    case R_TYPE_complex:
    case R_TYPE_character:
    case R_TYPE_raw:
        x = KEEP_N(r_vec_coerce(x, R_TYPE_list), &n_kept);
        // fallthrough
    case R_TYPE_list:
        check_unique_names(x);
        x = KEEP_N(r_list_as_environment(x, r_envs.empty), &n_kept);
        break;
    case R_TYPE_environment:
        break;
    default:
        r_abort(
            "`data` must be an uniquely named vector, list, data frame or "
            "environment"
        );
    }

    r_obj* pronoun = rlang_new_data_pronoun(x);

    FREE(n_kept);
    return pronoun;
}

static r_obj* data_mask_top_env_sym = NULL;

static void check_data_mask_input(r_obj* env, const char* arg) {
    if (r_typeof(env) != R_TYPE_environment) {
        r_abort(
            "Can't create data mask because `%s` must be an environment",
            arg
        );
    }
}
static void check_data_mask_top(r_obj* bottom, r_obj* top) {
    r_obj* cur = bottom;

    while (cur != r_envs.empty) {
        if (cur == top) {
            return;
        }
        cur = r_env_parent(cur);
    }

    r_abort("Can't create data mask because `top` is not a parent of `bottom`");
}

static r_obj* env_sym = NULL;
static r_obj* old_sym = NULL;
static r_obj* mask_sym = NULL;

static r_obj* tilde_fn = NULL;
static r_obj* restore_mask_fn = NULL;
static r_obj* restore_mask_formals = NULL;
static r_obj* restore_mask_body = NULL;

static void on_exit_restore_lexical_env(r_obj* mask, r_obj* old, r_obj* frame) {
    r_obj* env = KEEP(r_alloc_environment(2, r_envs.base));
    r_env_bind(env, mask_sym, mask);
    r_env_bind(env, old_sym, old);

    r_obj* fn =
        KEEP(r_new_function(restore_mask_formals, restore_mask_body, env));

    r_obj* call = KEEP(r_new_call(fn, r_null));
    r_on_exit(call, frame);

    FREE(3);
}

r_obj* ffi_new_data_mask(r_obj* bottom, r_obj* top) {
    r_obj* data_mask;

    if (bottom == r_null) {
        bottom = KEEP(r_alloc_environment(10, r_envs.empty));
        data_mask = bottom;
    } else {
        check_data_mask_input(bottom, "bottom");
        // Create a child because we don't know what might be in `bottom`
        // and we need to clear its contents without deleting any object
        // created in the data mask environment
        data_mask = KEEP(r_alloc_environment(10, bottom));
    }

    if (top == r_null) {
        top = bottom;
    } else {
        check_data_mask_input(top, "top");
    }
    if (top != bottom) {
        check_data_mask_top(bottom, top);
    }

    r_obj* ctxt_pronoun = KEEP(rlang_new_ctxt_pronoun(top));

    r_env_bind(data_mask, r_syms.tilde, tilde_fn);
    r_env_bind(data_mask, data_mask_flag_sym, data_mask);
    r_env_bind(data_mask, data_mask_env_sym, ctxt_pronoun);
    r_env_bind(data_mask, data_mask_top_env_sym, top);

    FREE(2);
    return data_mask;
}

r_obj* ffi_is_data_mask(r_obj* env) {
    return r_lgl(mask_info(env).type == RLANG_MASK_DATA);
}

static r_obj* mask_find(r_obj* env, r_obj* sym, bool* found) {
    if (r_typeof(sym) != R_TYPE_symbol) {
        r_abort("Internal error: Data pronoun must be subset with a symbol");
    }

    r_obj* top_env = r_env_has(env, data_mask_top_env_sym)
        ? r_env_get(env, data_mask_top_env_sym)
        : r_null;

    if (r_typeof(top_env) == R_TYPE_environment) {
        // Start lookup in the parent if the pronoun wraps a data mask
        env = r_env_parent(env);
    } else {
        // Data pronouns created from lists or data frames are converted
        // to a simple environment whose ancestry shouldn't be looked up.
        top_env = env;
    }
    int n_kept = 0;
    KEEP_N(top_env, &n_kept);

    r_obj* cur = env;
    do {
        if (r_env_has(cur, sym)) {
            *found = true;
            FREE(n_kept);
            return r_env_get(cur, sym);
        }

        if (cur == top_env) {
            break;
        } else {
            cur = r_env_parent(cur);
        }
    } while (cur != r_envs.empty);

    *found = false;
    FREE(n_kept);
    return r_null;
}
r_obj* ffi_data_pronoun_get(r_obj* pronoun, r_obj* sym, r_obj* error_call) {
    if (r_typeof(pronoun) != R_TYPE_environment) {
        r_abort("Internal error: Data pronoun must wrap an environment");
    }

    bool found;
    r_obj* obj = mask_find(pronoun, sym, &found);

    if (!found) {
        r_obj* call = KEEP(r_parse("abort_data_pronoun(x, call = y)"));
        r_eval_with_xy(call, sym, error_call, rlang_ns_env);
        r_abort("Internal error: .data subsetting should have failed earlier");
    }

    r_mark_shared(obj);
    return obj;
}

static void warn_env_as_mask_once(void) {
    const char* msg = "Passing an environment as data mask is deprecated.\n"
                      "Please use `new_data_mask()` to transform your "
                      "environment to a mask.\n"
                      "\n"
                      "  env <- env(foo = \"bar\")\n"
                      "\n"
                      "  # Bad:\n"
                      "  as_data_mask(env)\n"
                      "  eval_tidy(expr, env)\n"
                      "\n"
                      "  # Good:\n"
                      "  mask <- new_data_mask(env)\n"
                      "  eval_tidy(expr, mask)";
    deprecate_warn(msg, msg);
}

static r_obj* data_pronoun_sym = NULL;
static r_ssize mask_length(r_ssize n);

r_obj* ffi_as_data_mask(r_obj* data) {
    if (mask_info(data).type == RLANG_MASK_DATA) {
        return data;
    }
    if (data == r_null) {
        return ffi_new_data_mask(r_null, r_null);
    }

    int n_kept = 0;

    r_obj* bottom = NULL;

    switch (r_typeof(data)) {
    case R_TYPE_environment:
        warn_env_as_mask_once();
        bottom = KEEP_N(r_env_clone(data, NULL), &n_kept);
        break;

    case R_TYPE_logical:
    case R_TYPE_integer:
    case R_TYPE_double:
    case R_TYPE_complex:
    case R_TYPE_character:
    case R_TYPE_raw:
        data = r_vec_coerce(data, R_TYPE_list);
        KEEP_N(data, &n_kept);
        // fallthrough:

    case R_TYPE_list: {
        check_unique_names(data);

        r_obj* names = KEEP_N(r_names(data), &n_kept);

        r_ssize n_mask = mask_length(r_length(data));
        bottom = KEEP_N(r_alloc_environment(n_mask, r_envs.empty), &n_kept);

        if (names != r_null) {
            r_ssize n = r_length(data);

            r_obj* const* p_names = r_chr_cbegin(names);
            r_obj* const* p_data = r_list_cbegin(data);

            for (r_ssize i = 0; i < n; ++i) {
                // Ignore empty or missing names
                r_obj* nm = p_names[i];
                if (r_str_is_name(nm)) {
                    r_env_bind(bottom, r_str_as_symbol(nm), p_data[i]);
                }
            }
        }

        break;
    }

    default:
        r_abort("`data` must be a vector, list, data frame, or environment");
    }

    r_obj* data_mask = KEEP_N(ffi_new_data_mask(bottom, bottom), &n_kept);

    r_obj* data_pronoun = KEEP_N(ffi_as_data_pronoun(data_mask), &n_kept);
    r_env_bind(bottom, data_pronoun_sym, data_pronoun);

    FREE(n_kept);
    return data_mask;
}

static r_ssize mask_length(r_ssize n) {
    r_ssize n_grown =
        r_double_as_ssize(r_double_mult(r_ssize_as_double(n), 1.05));
    return r_ssize_max(n_grown, r_ssize_add(n, 20));
}

// For compatibility of the exported C callable
// TODO: warn
r_obj* ffi_new_data_mask_compat(r_obj* bottom, r_obj* top, r_obj* parent) {
    return ffi_new_data_mask(bottom, top);
}
r_obj* ffi_as_data_mask_compat(r_obj* data, r_obj* parent) {
    return ffi_as_data_mask(data);
}

static r_obj* tilde_prim = NULL;

static r_obj* base_tilde_eval(r_obj* tilde, r_obj* quo_env) {
    if (r_f_has_env(tilde)) {
        return tilde;
    }

    // Inline the base primitive because overscopes override `~` to make
    // quosures self-evaluate
    tilde = KEEP(r_new_call(tilde_prim, r_node_cdr(tilde)));
    tilde = KEEP(r_eval(tilde, quo_env));

    // Change it back because the result still has the primitive inlined
    r_node_poke_car(tilde, r_syms.tilde);

    FREE(2);
    return tilde;
}

r_obj* env_get_top_binding(r_obj* mask) {
    if (!r_env_has(mask, data_mask_top_env_sym)) {
        r_abort("Internal error: Can't find .top pronoun in data mask");
    }

    r_obj* top = r_env_get(mask, data_mask_top_env_sym);

    if (r_typeof(top) != R_TYPE_environment) {
        r_abort("Internal error: Unexpected .top pronoun type");
    }

    return top;
}

static r_obj* env_poke_parent_fn = NULL;
static r_obj* env_poke_fn = NULL;

r_obj* tilde_eval(r_obj* tilde, r_obj* current_frame, r_obj* caller_frame) {
    // Remove srcrefs from system call
    r_attrib_zap(tilde, r_syms.srcref);

    if (!is_quosure(tilde)) {
        return base_tilde_eval(tilde, caller_frame);
    }
    if (quo_is_missing(tilde)) {
        return r_missing_arg;
    }

    r_obj* expr = quo_get_expr(tilde);
    if (!r_is_symbolic(expr)) {
        return expr;
    }

    int n_kept = 0;

    r_obj* quo_env = KEEP_N(ffi_quo_get_env(tilde), &n_kept);
    if (r_typeof(quo_env) != R_TYPE_environment) {
        r_abort("Internal error: Quosure environment is corrupt");
    }

    r_obj* top = r_null;
    struct rlang_mask_info info = mask_info(caller_frame);

    switch (info.type) {
    case RLANG_MASK_DATA:
        top = KEEP_N(env_get_top_binding(info.mask), &n_kept);
        // Update `.env` pronoun to current quosure env temporarily
        poke_ctxt_env(info.mask, quo_env);
        break;
    case RLANG_MASK_QUOSURE:
        top = info.mask;
        break;
    case RLANG_MASK_NONE:
        r_abort("Internal error: Can't find the data mask");
    }

    // Unless the quosure was created in the mask, swap lexical contexts
    // temporarily by rechaining the top of the mask to the quosure
    // environment
    if (!r_env_inherits(quo_env, info.mask, top)) {
        // Unwind-protect the restoration of original parents
        on_exit_restore_lexical_env(
            info.mask,
            r_env_parent(top),
            current_frame
        );
        r_env_poke_parent(top, quo_env);
    }

    FREE(n_kept);
    return r_eval(expr, info.mask);
}

r_obj* ffi_tilde_eval(r_obj* call, r_obj* op, r_obj* args, r_obj* rho) {
    args = r_node_cdr(args);
    r_obj* tilde = r_node_car(args);
    args = r_node_cdr(args);
    r_obj* current_frame = r_node_car(args);
    args = r_node_cdr(args);
    r_obj* caller_frame = r_node_car(args);
    return tilde_eval(tilde, current_frame, caller_frame);
}

static const char* data_mask_objects_names[4] =
    {".__tidyeval_data_mask__.", "~", ".top_env", ".env"};

// Soft-deprecated in rlang 0.2.0
r_obj* ffi_data_mask_clean(r_obj* mask) {
    r_obj* bottom = r_env_parent(mask);
    r_obj* top = r_eval(data_mask_top_env_sym, mask);

    KEEP(top); // Help rchk

    if (top == r_null) {
        top = bottom;
    }

    // At this level we only want to remove our own stuff
    r_env_unbind_c_strings(
        mask,
        data_mask_objects_names,
        R_ARR_SIZEOF(data_mask_objects_names)
    );

    // Remove everything in the other levels
    r_obj* env = bottom;
    r_obj* parent = r_env_parent(top);
    while (env != parent) {
        r_obj* nms = KEEP(r_env_names(env));
        r_env_unbind_names(env, nms);
        FREE(1);
        env = r_env_parent(env);
    }

    FREE(1);
    return mask;
}

static r_obj* new_quosure_mask(r_obj* env) {
    r_obj* mask = KEEP(r_alloc_environment(3, env));
    r_env_bind(mask, r_syms.tilde, tilde_fn);
    r_env_bind(mask, quo_mask_flag_sym, mask);
    FREE(1);
    return mask;
}

r_obj* rlang_eval_tidy(r_obj* expr, r_obj* data, r_obj* env) {
    int n_kept = 0;

    if (is_quosure(expr)) {
        env = r_quo_get_env(expr);
        expr = r_quo_get_expr(expr);
    }

    // If there is no data, we only need to mask `~` with the definition
    // for quosure thunks. Otherwise we create a heavier data mask with
    // all the masking objects, data pronouns, etc.
    if (data == r_null) {
        r_obj* mask = KEEP_N(new_quosure_mask(env), &n_kept);
        r_obj* out = r_eval(expr, mask);
        FREE(n_kept);
        return out;
    }

    r_obj* mask = KEEP_N(ffi_as_data_mask(data), &n_kept);
    r_obj* top = KEEP_N(env_get_top_binding(mask), &n_kept);

    // Rechain the mask on the new lexical env but don't restore it on
    // exit. This way leaked masks inherit from a somewhat sensible
    // environment. We could do better with ALTENV and two-parent data
    // masks:
    //
    // * We'd create a new two-parents evaluation env for each quosure.
    //   The first parent would be the mask and the second the lexical
    //   environment.
    //
    // * The data mask top would always inherit from the empty
    //   environment.
    //
    // * Look-up in leaked environments would proceed from the data mask
    //   to the appropriate lexical environment (from quosures or from
    //   the `env` argument of eval_tidy()).
    if (!r_env_inherits(env, mask, top)) {
        poke_ctxt_env(mask, env);
        r_env_poke_parent(top, env);
    }

    r_obj* out = r_eval(expr, mask);
    FREE(n_kept);
    return out;
}

r_obj* ffi_eval_tidy(r_obj* call, r_obj* op, r_obj* args, r_obj* rho) {
    args = r_node_cdr(args);
    r_obj* expr = r_node_car(args);
    args = r_node_cdr(args);
    r_obj* data = r_node_car(args);
    args = r_node_cdr(args);
    r_obj* env = r_node_car(args);
    return rlang_eval_tidy(expr, data, env);
}

void rlang_init_eval_tidy(void) {
    r_obj* rlang_ns_env = KEEP(r_ns_env("rlang"));

    tilde_fn = r_eval(r_sym("tilde_eval"), rlang_ns_env);

    data_pronoun_class = r_chr("rlang_data_pronoun");
    r_preserve(data_pronoun_class);

    ctxt_pronoun_class = r_chr("rlang_ctxt_pronoun");
    r_preserve(ctxt_pronoun_class);

    empty_names_chr = r_alloc_character(2);
    r_preserve(empty_names_chr);
    r_chr_poke(empty_names_chr, 0, r_str(""));
    r_chr_poke(empty_names_chr, 1, r_globals.na_str);

    quo_mask_flag_sym = r_sym(".__tidyeval_quosure_mask__.");
    data_mask_flag_sym = r_sym(".__tidyeval_data_mask__.");
    data_mask_env_sym = r_sym(".env");
    data_mask_top_env_sym = r_sym(".top_env");
    data_pronoun_sym = r_sym(".data");

    tilde_prim = r_base_ns_get("~");
    env_poke_parent_fn = rlang_ns_get("env_poke_parent");
    env_poke_fn = rlang_ns_get("env_poke");

    env_sym = r_sym("env");
    old_sym = r_sym("old");
    mask_sym = r_sym("mask");

    restore_mask_fn = r_parse_eval(
        "function() {                          \n"
        "  ctxt_pronoun <- `mask`$.env         \n"
        "  if (!is.null(ctxt_pronoun)) {       \n"
        "    parent.env(ctxt_pronoun) <- `old` \n"
        "  }                                   \n"
        "                                      \n"
        "  top <- `mask`$.top_env              \n"
        "  if (is.null(top)) {                 \n"
        "    top <- `mask`                     \n"
        "  }                                   \n"
        "                                      \n"
        "  parent.env(top) <- `old`            \n"
        "}                                     \n",
        r_envs.base
    );
    r_preserve(restore_mask_fn);

    restore_mask_formals = r_fn_formals(restore_mask_fn);
    r_preserve(restore_mask_formals);

    restore_mask_body = r_fn_body(restore_mask_fn);
    r_preserve(restore_mask_body);

    FREE(1);
}
