#include <assert.h>

#include "apfl.h"
#include "hashmap.h"
#include "internal.h"
#include "resizable.h"

struct apfl_ctx_data {
    apfl_hashmap scope;
};

struct variable_data {
    unsigned refcount;
    struct apfl_value value;
};

typedef struct variable_data *variable;

static struct apfl_result evaluate(apfl_ctx, struct apfl_expr *);
static variable variable_new(void);
static variable variable_incref(variable var);
static void variable_unref(variable var);

static bool
scope_keys_eq(void *opaque, const void *_a, const void *_b)
{
    (void)opaque;
    apfl_refcounted_string const *a = _a;
    apfl_refcounted_string const *b = _b;

    return *a == *b || apfl_string_eq(*a, *b);
}

static apfl_hash
scope_calc_hash(void *opaque, const void *_key)
{
    (void)opaque;
    apfl_refcounted_string const *key = _key;
    struct apfl_string_view sv = apfl_string_view_from(*key);
    return apfl_hash_fnv1a(sv.bytes, sv.len);
}

static void
scope_destroy_key(void *opaque, void *_key)
{
    (void)opaque;
    apfl_refcounted_string *key = _key;

    apfl_refcounted_string_unref(*key);
}

static void
scope_destroy_value(void *opaque, void *_value)
{
    (void)opaque;
    variable *value = _value;
    variable_unref(*value);
}

static void
scope_copy_key(void *opaque, void *_dest, void *_src)
{
    (void)opaque;
    apfl_refcounted_string *dest = _dest;
    apfl_refcounted_string *src = _src;
    *dest = apfl_refcounted_string_incref(*src);
}

static void
scope_copy_value(void *opaque, void *_dest, void *_src)
{
    (void)opaque;
    variable *dest = _dest;
    variable *src = _src;
    *dest = variable_incref(*src);
}

static const struct apfl_hashmap_callbacks scope_hashmap_callbacks = {
    .opaque        = NULL,
    .keys_eq       = scope_keys_eq,
    .calc_hash     = scope_calc_hash,
    .destroy_key   = scope_destroy_key,
    .destroy_value = scope_destroy_value,
    .copy_key      = scope_copy_key,
    .copy_value    = scope_copy_value,
};

apfl_hashmap
scope_new(void)
{
    return apfl_hashmap_new(
        scope_hashmap_callbacks,
        sizeof(apfl_refcounted_string),
        sizeof(variable)
    );
}

static variable
variable_new(void)
{
    variable var = ALLOC(struct variable_data);
    if (var == NULL) {
        return NULL;
    }

    var->refcount = 1;
    var->value.type = APFL_VALUE_NIL;

    return var;
}

static variable
variable_incref(variable var)
{
    if (var != NULL) {
        var->refcount++;
    }
    return var;
}

static void
variable_unref(variable var)
{
    if (var != NULL && apfl_refcount_dec(&var->refcount)) {
        apfl_value_deinit(&var->value);
        free(var);
    }
}

apfl_ctx
apfl_ctx_new(void)
{
    apfl_ctx ctx = ALLOC(struct apfl_ctx_data);
    if (ctx == NULL) {
        return NULL;
    }

    if ((ctx->scope = scope_new()) == NULL) {
        free(ctx);
        return NULL;
    }

    return ctx;
}

void
apfl_ctx_destroy(apfl_ctx ctx)
{
    if (ctx == NULL) {
        return;
    }

    apfl_hashmap_destroy(ctx->scope);
    free(ctx);
}

static variable
ctx_get_var_for_assignment_inner(apfl_ctx ctx, /*borrowed*/ apfl_refcounted_string name)
{
    variable var;
    if (apfl_hashmap_get(ctx->scope, &name, &var)) {
        return var;
    }

    if ((var = variable_new()) == NULL) {
        return NULL;
    }

    if (!apfl_hashmap_set(ctx->scope, &name, &var)) {
        variable_unref(var);
        return NULL;
    }

    return var;
}

static variable
ctx_get_var_for_assignment(apfl_ctx ctx, apfl_refcounted_string name)
{
    variable var = ctx_get_var_for_assignment_inner(ctx, name);
    apfl_refcounted_string_unref(name);
    return var;
}

static variable
ctx_get_var(apfl_ctx ctx, apfl_refcounted_string name)
{
    variable var;
    bool ok = apfl_hashmap_get(ctx->scope, &name, &var);
    apfl_refcounted_string_unref(name);
    return ok ? var : NULL;
}

static struct apfl_result
fatal(void)
{
    return (struct apfl_result) { .type = APFL_RESULT_ERR_FATAL };
}

static struct apfl_result
evaluate_constant(struct apfl_expr_const *constant)
{
    apfl_refcounted_string rcstring;

    switch (constant->type) {
    case APFL_EXPR_CONST_NIL:
        return (struct apfl_result) {
            .type = APFL_RESULT_OK,
            .value = {
                .type = APFL_VALUE_NIL,
            },
        };
    case APFL_EXPR_CONST_BOOLEAN:
        return (struct apfl_result) {
            .type = APFL_RESULT_OK,
            .value = {
                .type = APFL_VALUE_BOOLEAN,
                .boolean = constant->boolean,
            },
        };
    case APFL_EXPR_CONST_STRING:
        rcstring = apfl_string_move_into_new_refcounted(&constant->string);
        if (rcstring == NULL) {
            return fatal();
        }

        return (struct apfl_result) {
            .type = APFL_RESULT_OK,
            .value = {
                .type = APFL_VALUE_STRING,
                .string = rcstring,
            },
        };
    case APFL_EXPR_CONST_NUMBER:
        return (struct apfl_result) {
            .type = APFL_RESULT_OK,
            .value = {
                .type = APFL_VALUE_NUMBER,
                .number = constant->number,
            },
        };
    }

    assert(false);
}

static struct apfl_result
evaluate_list_expr_to_list(apfl_ctx ctx, struct apfl_expr *expr, apfl_editable_list elist)
{
    struct apfl_result result = evaluate(ctx, expr);
    if (result.type != APFL_RESULT_OK) {
        return result;
    }

    if (result.value.type != APFL_VALUE_LIST) {
        apfl_value_deinit(&result.value);
        return (struct apfl_result) { .type = APFL_RESULT_ERR };
    }

    apfl_list list = apfl_list_incref(result.value.list);
    apfl_value_deinit(&result.value);

    if (!apfl_editable_list_append_list(elist, list)) {
        return fatal();
    }

    return (struct apfl_result) { .type = APFL_RESULT_OK, .value.type = APFL_VALUE_NIL };
}

static struct apfl_result
evaluate_expr_list_into_list(apfl_ctx ctx, struct apfl_expr_list *list, apfl_editable_list elist)
{
    for (size_t i = 0; i < list->len; i++) {
        struct apfl_expr_list_item *item = &list->items[i];

        if (item->expand) {
            struct apfl_result result = evaluate_list_expr_to_list(ctx, item->expr, elist);
            if (result.type != APFL_RESULT_OK) {
                return result;
            }
        } else {
            struct apfl_result result = evaluate(ctx, item->expr);
            if (result.type != APFL_RESULT_OK) {
                return result;
            }

            if (!apfl_editable_list_append(elist, apfl_value_move(&result.value))) {
                return fatal();
            }
        }
    }

    return (struct apfl_result) { .type = APFL_RESULT_OK, .value.type = APFL_VALUE_NIL };
}

static struct apfl_result
evaluate_list(apfl_ctx ctx, struct apfl_expr_list *list)
{
    apfl_editable_list elist = apfl_editable_list_new();
    if (elist == NULL) {
        return fatal();
    }

    struct apfl_result result = evaluate_expr_list_into_list(ctx, list, elist);
    if (result.type != APFL_RESULT_OK) {
        apfl_editable_list_destroy(elist);
        return result;
    }

    apfl_list out = apfl_editable_list_finalize(elist);
    if (out == NULL) {
        return fatal();
    }

    return (struct apfl_result) {
        .type = APFL_RESULT_OK,
        .value = {
            .type = APFL_VALUE_LIST,
            .list = out,
        },
    };
}

static struct apfl_result
evaluate_dict(apfl_ctx ctx, struct apfl_expr_dict *dict)
{
    apfl_editable_dict ed = apfl_editable_dict_new();

    if (ed == NULL) {
        return fatal();
    }

    for (size_t i = 0; i < dict->len; i++) {
        struct apfl_result result;
        struct apfl_value key, value;

        result = evaluate(ctx, dict->items[i].k);
        if (result.type != APFL_RESULT_OK) {
            apfl_editable_dict_destroy(ed);
            return result;
        }
        key = result.value;

        result = evaluate(ctx, dict->items[i].v);
        if (result.type != APFL_RESULT_OK) {
            apfl_editable_dict_destroy(ed);
            apfl_value_deinit(&key);
            return result;
        }
        value = result.value;

        if (!apfl_editable_dict_set(
            ed,
            apfl_value_move(&key),
            apfl_value_move(&value)
        )) {
            apfl_editable_dict_destroy(ed);
            return fatal();
        }
    }

    apfl_dict out = apfl_editable_dict_finalize(ed);
    if (out == NULL) {
        return fatal();
    }

    return (struct apfl_result) {
        .type = APFL_RESULT_OK,
        .value = {
            .type = APFL_VALUE_DICT,
            .dict = out,
        },
    };
}

static struct apfl_result
evaluate_dot(apfl_ctx ctx, struct apfl_expr_dot *dot)
{
    struct apfl_result result;

    result = evaluate(ctx, dot->lhs);
    if (result.type != APFL_RESULT_OK) {
        return result;
    }
    struct apfl_value lhs = result.value;

    apfl_refcounted_string rcstring = apfl_string_move_into_new_refcounted(&dot->rhs);
    if (rcstring == NULL) {
        return fatal();
    }
    struct apfl_value key = (struct apfl_value) {
        .type = APFL_VALUE_STRING,
        .string = rcstring,
    };

    struct apfl_value out;
    if (apfl_value_get_item(apfl_value_move(&lhs), apfl_value_move(&key), &out) == APFL_VALUE_GET_ITEM_OK) {
        return (struct apfl_result) {
            .type = APFL_RESULT_OK,
            .value = out,
        };
    } else {
        return (struct apfl_result) { .type = APFL_RESULT_ERR }; // TODO: Describe error
    }
}

static struct apfl_result
evaluate_at(apfl_ctx ctx, struct apfl_expr_at *at)
{
    struct apfl_result result;

    result = evaluate(ctx, at->lhs);
    if (result.type != APFL_RESULT_OK) {
        return result;
    }
    struct apfl_value lhs = result.value;

    result = evaluate(ctx, at->rhs);
    if (result.type != APFL_RESULT_OK) {
        apfl_value_deinit(&lhs);
        return result;
    }
    struct apfl_value rhs = result.value;

    struct apfl_value out;
    if (apfl_value_get_item(apfl_value_move(&lhs), apfl_value_move(&rhs), &out) == APFL_VALUE_GET_ITEM_OK) {
        return (struct apfl_result) {
            .type = APFL_RESULT_OK,
            .value = out,
        };
    } else {
        return (struct apfl_result) { .type = APFL_RESULT_ERR }; // TODO: Describe error
    }
}

static struct apfl_result
evaluate_assignment(apfl_ctx ctx, struct apfl_expr_assignment *assignment)
{
    // TODO: Use assignment->local flag. Sonce we don't hev functions yet, it's not yet relevant.

    if (
        assignment->lhs.type != APFL_EXPR_ASSIGNABLE_VAR_OR_MEMBER
        || assignment->lhs.var_or_member.type != APFL_EXPR_ASSIGNABLE_VAR_OR_MEMBER_VAR
    ) {
        return (struct apfl_result) { .type = APFL_RESULT_ERR }; // TODO: Implement other assignable types
    }

    apfl_refcounted_string varname = apfl_string_move_into_new_refcounted(&assignment->lhs.var_or_member.var);
    if (varname == NULL) {
        return fatal();
    }

    variable var = ctx_get_var_for_assignment(ctx, varname);

    if (var == NULL) {
        return fatal();
    }

    struct apfl_result result = evaluate(ctx, assignment->rhs);
    if (result.type != APFL_RESULT_OK) {
        variable_unref(var);
        return result;
    }

    apfl_value_deinit(&var->value);
    var->value = apfl_value_incref(result.value);
    variable_unref(var);
    return result;
}

static struct apfl_result
evaluate_var(apfl_ctx ctx, struct apfl_string *varname)
{
    apfl_refcounted_string rcstring = apfl_string_move_into_new_refcounted(varname);
    if (rcstring == NULL) {
        return fatal();
    }

    variable var = ctx_get_var(ctx, rcstring);
    if (var == NULL) {
        return (struct apfl_result) { .type = APFL_RESULT_ERR };
    }

    struct apfl_value value = apfl_value_incref(var->value);
    variable_unref(var);
    return (struct apfl_result) {
        .type = APFL_RESULT_OK,
        .value = value,
    };
}

static struct apfl_result
evaluate(apfl_ctx ctx, struct apfl_expr *expr)
{
    switch (expr->type) {
    case APFL_EXPR_CONSTANT:
        return evaluate_constant(&expr->constant);
    case APFL_EXPR_LIST:
        return evaluate_list(ctx, &expr->list);
    case APFL_EXPR_DICT:
        return evaluate_dict(ctx, &expr->dict);
    case APFL_EXPR_DOT:
        return evaluate_dot(ctx, &expr->dot);
    case APFL_EXPR_AT:
        return evaluate_at(ctx, &expr->at);
    case APFL_EXPR_ASSIGNMENT:
        return evaluate_assignment(ctx, &expr->assignment);
    case APFL_EXPR_VAR:
        return evaluate_var(ctx, &expr->var);
    case APFL_EXPR_BLANK:
        return (struct apfl_result) {
            .type = APFL_RESULT_OK,
            .value = {
                .type = APFL_VALUE_NIL,
            },
        };
    case APFL_EXPR_CALL:
    case APFL_EXPR_SIMPLE_FUNC:
    case APFL_EXPR_COMPLEX_FUNC:
        break; // Not implemented yet
    }

    return (struct apfl_result) { .type = APFL_RESULT_ERR };
}

struct apfl_result
apfl_eval(apfl_ctx ctx, struct apfl_expr expr)
{
    struct apfl_result result = evaluate(ctx, &expr);
    apfl_expr_deinit(&expr);
    return result;
}