#include <rlang.h>
#include "internal.h"
#include "vec.h"

#include "decl/attr-decl.h"

r_obj* ffi_names2(r_obj* x, r_obj* env) {
    const enum r_type type = r_typeof(x);

    if (type == R_TYPE_environment) {
        r_abort("Use `env_names()` for environments.");
    }

    // Handle pairlists and language objects specially like `getAttrib()`
    // does. `r_names()` will not find these names because it has a guarantee
    // to never allocate.
    if (type == R_TYPE_pairlist || type == R_TYPE_call) {
        return node_names(x);
    }

    r_obj* nms;
    if (r_is_object(x)) {
        nms = KEEP(names_dispatch(x, env));
    } else {
        nms = KEEP(r_names(x));
    }

    if (nms == r_null) {
        r_ssize n = r_length(x);
        nms = KEEP(r_alloc_character(n));
        r_chr_fill(nms, r_strs.empty, n);
    } else {
        nms = KEEP(ffi_replace_na(nms, r_chrs.empty_string));
    }

    FREE(2);
    return nms;
}

static r_obj* node_names(r_obj* x) {
    r_ssize n = r_length(x);

    r_obj* out = KEEP(r_alloc_character(n));

    int i = 0;

    for (; x != r_null; x = r_node_cdr(x), ++i) {
        r_obj* tag = r_node_tag(x);

        if (tag == r_null) {
            r_chr_poke(out, i, r_strs.empty);
        } else {
            r_chr_poke(out, i, PRINTNAME(tag));
        }
    }

    FREE(1);
    return out;
}

r_obj* ffi_set_names(r_obj* x, r_obj* mold, r_obj* nm, r_obj* env) {
    int n_kept = 0;

    r_obj* dots = KEEP_N(rlang_dots(env), &n_kept);

    if (!r_is_vector(x, -1)) {
        r_abort("`x` must be a vector");
    }

    if (nm == r_null) {
        x = set_names_dispatch(x, r_null, env);

        FREE(n_kept);
        return x;
    }

    if (r_is_function(nm) || r_is_formula(nm, -1, -1)) {
        if (r_names(mold) == r_null) {
            mold = KEEP_N(eval_as_character(mold, env), &n_kept);
        } else {
            mold = KEEP_N(ffi_names2(mold, env), &n_kept);
        }

        nm = KEEP_N(rlang_as_function(nm, env), &n_kept);
        nm = KEEP_N(eval_fn_dots(nm, mold, dots, env), &n_kept);
    } else {
        if (r_length(dots) > 0) {
            nm = KEEP_N(eval_fn_dots(c_fn, nm, dots, env), &n_kept);
        }

        nm = KEEP_N(eval_as_character(nm, env), &n_kept);
    }

    r_ssize n;
    if (r_is_object(x)) {
        n = length_dispatch(x, env);
    } else {
        n = r_length(x);
    }

    if (r_typeof(nm) != R_TYPE_character) {
        r_abort("`nm` must be `NULL` or a character vector.");
    }

    r_ssize nm_n = r_length(nm);
    if (nm_n != n) {
        if (nm_n != 1) {
            r_abort(
                "The size of `nm` (%d) must be compatible with the size of `x` "
                "(%d).",
                nm_n,
                n
            );
        }

        // Recycle names vector of size 1.
        // TODO: ALTREP repetitions?
        r_obj* val = r_chr_get(nm, 0);
        nm = KEEP_N(r_alloc_character(n), &n_kept);
        r_chr_fill(nm, val, n);
    }

    if (!is_character(nm, n, OPTION_BOOL_null, OPTION_BOOL_null)) {
        r_abort(
            "`nm` must be `NULL` or a character vector the same length as `x`"
        );
    }

    x = set_names_dispatch(x, nm, env);

    FREE(n_kept);
    return x;
}

static r_obj* eval_fn_dots(r_obj* fn, r_obj* x, r_obj* dots, r_obj* env) {
    r_obj* args = KEEP(r_new_node(r_syms.dot_x, dots));
    r_obj* call = KEEP(r_new_call(r_syms.dot_fn, args));

    // This evaluates `.fn(.x, ...)`
    // `.x` is the first input, x
    // `.fn` is the function, fn
    // The dots are a pairlist already in the call
    r_env_bind(env, r_syms.dot_x, x);
    r_env_bind(env, r_syms.dot_fn, fn);

    r_obj* out = r_eval(call, env);

    FREE(2);
    return out;
}

static inline r_obj* eval_as_character(r_obj* x, r_obj* env) {
    r_env_bind(env, r_syms.dot_x, x);
    return r_eval(as_character_call, env);
}

static inline r_obj* names_dispatch(r_obj* x, r_obj* env) {
    r_env_bind(env, r_syms.dot_x, x);
    return r_eval(names_call, env);
}

// Use `names<-()` rather than setting names directly with
// `r_attrib_poke_names()` for genericity and for speed. `names<-()` can shallow
// duplicate `x`'s attributes using ALTREP wrappers, which is not in R's public
// API.
static inline r_obj* set_names_dispatch(r_obj* x, r_obj* nm, r_obj* env) {
    r_env_bind(env, r_syms.dot_x, x);
    r_env_bind(env, r_syms.dot_y, nm);
    return r_eval(set_names_call, env);
}

static inline r_ssize length_dispatch(r_obj* x, r_obj* env) {
    r_env_bind(env, r_syms.dot_x, x);
    r_obj* n = KEEP(r_eval(length_call, env));

    if (r_length(n) != 1) {
        r_abort("Object length must have size 1, not %i", r_length(n));
    }

    r_ssize out;

    switch (r_typeof(n)) {
    case R_TYPE_integer:
        out = (r_ssize) r_int_begin(n)[0];
        break;
    case R_TYPE_double:
        out = r_dbl_begin(n)[0];
        break;
    default:
        r_abort(
            "Object length has unknown type %s",
            r_type_as_c_string(r_typeof(n))
        );
    }

    FREE(1);
    return out;
}

r_obj* zap_srcref(r_obj* x) {
    switch (r_typeof(x)) {
    case R_TYPE_call:
        return call_zap_srcref(x);
    case R_TYPE_closure:
        return fn_zap_srcref(x);
    case R_TYPE_expression:
        return expr_vec_zap_srcref(x);
    default:
        return x;
    }
}

static r_obj* fn_zap_srcref(r_obj* x) {
    r_obj* formals = r_fn_formals(x);
    r_obj* body = r_fn_body(x);
    r_obj* env = r_fn_env(x);

    body = KEEP(zap_srcref(body));

    r_obj* out = KEEP(r_new_function(formals, body, env));

    r_attrib_clone_from(out, x);
    r_attrib_zap(out, r_syms.srcref);

    FREE(2);
    return out;
}

static r_obj* call_zap_srcref(r_obj* x) {
    x = KEEP(r_clone(x));

    attrib_zap_srcref(x);

    if (r_node_car(x) == r_syms.function) {
        // Remove `call[[4]]` where the parser stores srcref information
        // for calls to `function`
        r_node_poke_cdr(r_node_cddr(x), r_null);
    }

    for (r_obj* node = x; node != r_null; node = r_node_cdr(node)) {
        r_node_poke_car(node, zap_srcref(r_node_car(node)));
    }

    FREE(1);
    return x;
}

static r_obj* expr_vec_zap_srcref(r_obj* x) {
    x = KEEP(r_clone(x));

    attrib_zap_srcref(x);

    // `VECTOR_PTR_RO()` (i.e. `r_list_cbegin()`) doesn't allow `EXPRSXP`, so we
    // use rlang equivalents of `Rf_xlength()`, `VECTOR_ELT()`, and
    // `SET_VECTOR_ELT()` instead
    r_ssize n = r_length(x);

    for (r_ssize i = 0; i < n; ++i) {
        r_obj* elt = r_list_get(x, i);
        r_list_poke(x, i, zap_srcref(elt));
    }

    FREE(1);
    return x;
}

static void attrib_zap_srcref(r_obj* x) {
    r_attrib_zap(x, r_syms.srcfile);
    r_attrib_zap(x, r_syms.srcref);
    r_attrib_zap(x, r_syms.wholeSrcref);
}

void rlang_init_attr(r_obj* ns) {
    c_fn = r_eval(r_sym("c"), r_envs.base);

    as_character_call = r_parse("as.character(.x)");
    r_preserve(as_character_call);

    names_call = r_parse("names(.x)");
    r_preserve(names_call);

    set_names_call = r_parse("`names<-`(.x, .y)");
    r_preserve(set_names_call);

    length_call = r_parse("length(.x)");
    r_preserve(length_call);
}

static r_obj* c_fn = NULL;
static r_obj* as_character_call = NULL;
static r_obj* names_call = NULL;
static r_obj* set_names_call = NULL;
static r_obj* length_call = NULL;
