#include "preprocess.h"

#include <assert.h>

#include "built-ins.h"
#include "func.h"
#include "value.h"

void ypreprocess(yexpr_t* expr) {
    ypreprocess_resolve_builtins(expr);
    ypreprocess_refs(expr);
}

void ypreprocess_resolve_builtins(yexpr_t* expr) {
    switch (expr->typ) {
        case YEXPR_NUMBER:
        case YEXPR_STRING:
        case YEXPR_ATOM:
        case YEXPR_EXCEPTION:
        case YEXPR_UNDEF:
            break;
        case YEXPR_BUILTIN:
            assert(expr->typ != YEXPR_BUILTIN /* We shouldn't traverse the same expression twice */);
        case YEXPR_REF:
            if (yref_type(&expr->value.ref) == YREF_SYM)
                ybuiltin_translate(expr);
            break;
        case YEXPR_LIST:
            if (expr->value.list.len == 0) break;
            yvec_t* list = &expr->value.list;
            for (size_t i = 0; i < list->len; i++) {
                ypreprocess_resolve_builtins(YVEC_AT(list, i, yexpr_t));
            }
            break;
        default:
            assert(false /* unexpected expression type! */);
    }
}

/**
 * @brief Resolve references to arguments in the function body.
 */
static void ypreprocess_resolve_inner_refs(yexpr_t* body, yvec_t* arg_descs) {
    switch (body->typ) {
        case YEXPR_REF:
            if (yref_type(&body->value.ref) != YREF_SYM) break;
            // Check if reference is to known argument.
            for (size_t i = 0; i < arg_descs->len; i++) {
                yarg_desc_t* arg = YVEC_AT(arg_descs, i, yarg_desc_t);
                if (0 == ystr_cmp(&body->value.ref.ref.sym, &arg->argname)) {
                    yexpr_set_ref(body, yref_clone(&arg->argref));
                    break;
                }
            }
            break;
        case YEXPR_LIST:
            for (size_t i = 0; i < body->value.list.len; i++) {
                yexpr_t* expr = YVEC_AT(&body->value.list, i, yexpr_t);
                ypreprocess_resolve_inner_refs(expr, arg_descs);
            }
            break;
        default:
            break;
    }
}

/**
 * @brief Store a `(defn ...)` as function and replace it with a YEXPR_UNDEF
 * definition.
 *
 * `expr` must be a list starting with the `defn` built-in.
 *
 * Returns false if syntax is incorrect. The reference to the function is stored
 * into `newref`.
 */
static bool ypreprocess_compile_defn(yexpr_t* expr, yref_t* newref) {
    yvec_t* list = &expr->value.list;
    assert(4 <= list->len);
    yexpr_t* defn = YVEC_AT(list, 0, yexpr_t);
    yexpr_t* name = YVEC_AT(list, 1, yexpr_t);
    yexpr_t* args = YVEC_AT(list, 2, yexpr_t);
    if (defn->typ != YEXPR_BUILTIN || defn->value.builtin != YBUILTIN_DEFN)
        return false;
    if (name->typ != YEXPR_REF) return false;
    assert(yref_type(&name->value.ref) == YREF_SYM);
    if (args->typ != YEXPR_LIST) return false;
    size_t n_args = args->value.list.len;

    yfunc_t func;
    yvec_t body_exprs = YVEC_NEW(yvec_at(list, 3), list->len - 3, yexpr_t);
    yvec_t arg_descs = YVEC_NEW(NULL, args->value.list.len, yarg_desc_t);
    for (size_t i = 0; i < args->value.list.len; i++) {
        yarg_desc_t desc;
        yexpr_t* arg = YVEC_AT(&args->value.list, i, yexpr_t);
        assert(arg->typ == YEXPR_REF && yref_type(&arg->value.ref) == YREF_SYM);
        desc.argname = ystr_clone(&arg->value.ref.ref.sym);
        desc.argref = yref_new_cix(n_args - 1 - i);
        YVEC_PUSH(&arg_descs, &desc);
    }

    func.name = name->value.ref.ref.sym;
    func.args = arg_descs;
    yexpr_t body = yexpr_new();
    yexpr_set_list(&body, body_exprs);  // no destroy!
    // Resolve references to arguments in body.
    ypreprocess_resolve_inner_refs(&body, &arg_descs);
    func.body = body_exprs;

    // Store function as value at the named reference.
    yvalue_t func_value;
    func_value.typ = YVALUE_FUNC;
    func_value.value.func = func;

    *newref = yvalue_set(YVALUE_INSERT, &func_value);

    yexpr_destroy(defn);
    yexpr_destroy(args);
    yvec_destroy(list);

    *expr = yexpr_new();
    return true;
}

typedef struct {
    ystr_t sym;
    yref_t ref;
} ypreprocess_existing_t;

/// References in a scope.
typedef struct {
    /// Vector of ypreprocess_existing_t
    yvec_t refs;
} ypreprocess_scope_t;

static yref_t ypreprocess_find_or_create_ref(
    ystr_t* symbol, yvec_t* /* of ypreprocess_scope_t* */ scope_stack) {
    // Search for symbol in stack, starting at top-most scope (enclosing list)
    // and working upwards.
    for (size_t i = 0; i < scope_stack->len; i++) {
        size_t index = scope_stack->len - 1 - i;
        ypreprocess_scope_t* scope =
            *YVEC_AT(scope_stack, index, ypreprocess_scope_t*);

        for (size_t ref = 0; ref < scope->refs.len; ref++) {
            ypreprocess_existing_t* candidate =
                YVEC_AT(&scope->refs, ref, ypreprocess_existing_t);
            if (0 == ystr_cmp(symbol, &candidate->sym)) {
                return candidate->ref;
            }
        }
    }

    // No existing reference found! Create one in the current scope.
    yref_t ref = yref_new();
    // For all new references, we define them in the scope above the current
    // list (thus "- 2"), e.g. in let, defn
    ypreprocess_scope_t* topscope =
        *YVEC_AT(scope_stack, scope_stack->len - 2, ypreprocess_scope_t*);
    ypreprocess_existing_t newref = {.sym = ystr_clone(symbol), .ref = ref};
    YVEC_PUSH(&topscope->refs, &newref);
    return ref;
}

static void ypreprocess_refs_recursive(
    yvec_t* /* of ypreprocess_scope_t* */ scope_stack, yexpr_t* expr) {
    // Abort cheaply for types we don't want to resolve anyway.
    switch (expr->typ) {
        case YEXPR_UNDEF:
        case YEXPR_BUILTIN:
        case YEXPR_NUMBER:
        case YEXPR_STRING:
        case YEXPR_ATOM:
        case YEXPR_EXCEPTION:
            return;
        default:
            break;
    }

    ypreprocess_scope_t scope;
    ypreprocess_scope_t* scopep = &scope;

    switch (expr->typ) {
        case YEXPR_REF:
            if (yref_type(&expr->value.ref) == YREF_SYM) {
                ystr_t* sym = &expr->value.ref.ref.sym;
                yref_t resolved =
                    ypreprocess_find_or_create_ref(sym, scope_stack);
                yexpr_set_ref(expr, resolved);
            }
            break;
        case YEXPR_LIST:
            // Create new scope for list.
            scope.refs = YVEC_NEW(NULL, 4, ypreprocess_existing_t);
            YVEC_PUSH(scope_stack, &scopep);

            // Handle functions.
            if (expr->value.list.len >= 4) {
                yexpr_t* head = YVEC_AT(&expr->value.list, 0, yexpr_t);
                if (head->typ == YEXPR_BUILTIN &&
                    head->value.builtin == YBUILTIN_DEFN) {
                    // Compile and store function. Push symbol->reference
                    // mapping onto scope stack.
                    yref_t new_ref;
                    // TODO: handle errors better
                    assert(ypreprocess_compile_defn(expr, &new_ref));
                    yvalue_t* func = yvalue_get(new_ref);
                    assert(func != NULL);

                    ypreprocess_existing_t funcref = {
                        .sym = ystr_clone(&func->value.func.name),
                        .ref = new_ref};
                    YVEC_PUSH(&(*YVEC_AT(scope_stack, scope_stack->len - 2,
                                         ypreprocess_scope_t*))
                                   ->refs,
                              &funcref);

                    yvec_t* func_body = &func->value.func.body;
                    for (size_t i = 0; i < func_body->len; i++) {
                        ypreprocess_refs_recursive(
                            scope_stack, YVEC_AT(func_body, i, yexpr_t));
                    }

                    goto cleanup;
                }
            }

            // Resolve list elements recursively.
            for (size_t i = 0; i < expr->value.list.len; i++) {
                yexpr_t* elem = YVEC_AT(&expr->value.list, i, yexpr_t);
                ypreprocess_refs_recursive(scope_stack, elem);
            }

        cleanup:
            // Clean up.
            yvec_pop(scope_stack, NULL);
            for (size_t i = 0; i < scope.refs.len; i++) {
                ypreprocess_existing_t* ex =
                    YVEC_AT(&scope.refs, i, ypreprocess_existing_t);
                ystr_destroy(&ex->sym);
                yref_destroy(&ex->ref);
            }
            yvec_destroy(&scope.refs);
            break;
        default:
            break;
    }
}

// Initial invocation must be of type YEXPR_LIST.
void ypreprocess_refs(yexpr_t* expr) {
    yvec_t scope_stack = YVEC_NEW(NULL, 16, ypreprocess_scope_t*);
    ypreprocess_refs_recursive(&scope_stack, expr);
    assert(scope_stack.len == 0);
    yvec_destroy(&scope_stack);
}

void ypreprocess_refs_repl(yexpr_t* repl_expr, yvec_t** scope_stack) {
    if (*scope_stack == NULL) {
        *scope_stack = malloc(sizeof(yvec_t));
        **scope_stack = YVEC_NEW(NULL, 16, ypreprocess_scope_t*);
        ypreprocess_scope_t* root = malloc(sizeof(ypreprocess_scope_t));
        root->refs = YVEC_NEW(NULL, 4, ypreprocess_existing_t);
        YVEC_PUSH(*scope_stack, &root);
    }
    assert(YEXPR_LIST == repl_expr->typ);
    for (size_t i = 0; i < repl_expr->value.list.len; i++) {
        yexpr_t* elem = YVEC_AT(&repl_expr->value.list, i, yexpr_t);
        ypreprocess_refs_recursive(*scope_stack, elem);
    }
}