#include <assert.h>

#include "apfl.h"

#include "alloc.h"
#include "bytecode.h"
#include "compile.h"
#include "context.h"
#include "format.h"
#include "hashmap.h"
#include "strings.h"
#include "value.h"

#define TRY(ex)                         \
    do {                                \
        enum apfl_result result = (ex); \
        if (result != APFL_RESULT_OK) { \
            return result;              \
        }                               \
    } while (0)

static enum apfl_result
stack_push_or_fatal(apfl_ctx ctx, struct apfl_value value)
{
    return apfl_stack_push(ctx, value) ? APFL_RESULT_OK : APFL_RESULT_ERR_FATAL;
}

static void
stack_must_drop(apfl_ctx ctx, apfl_stackidx index)
{
    assert(apfl_stack_drop(ctx, index));
}

static enum apfl_result
get_argument(size_t *i, struct instruction_list *ilist, union instruction_or_arg *arg)
{
    if (*i >= ilist->len) {
        return APFL_RESULT_ERR;
    }

    *arg = ilist->instructions[(*i)++];
    return APFL_RESULT_OK;
}

static enum apfl_result
variable_get(apfl_ctx ctx, struct apfl_string *name)
{
    struct apfl_value value;
    if (!apfl_scope_get(ctx->scope, name, &value)) {
        return APFL_RESULT_ERR;
    }

    return stack_push_or_fatal(ctx, value);
}

static enum apfl_result
variable_set(apfl_ctx ctx, struct apfl_string *name, bool keep_on_stack)
{
    struct apfl_value value;
    if (!apfl_stack_get(ctx, &value, -1)) {
        return APFL_RESULT_ERR;
    }
    if (!apfl_scope_set(&ctx->gc, ctx->scope, name, value)) {
        return APFL_RESULT_ERR_FATAL;
    }
    if (keep_on_stack) {
        // If the value should be kept on the stack, the value is now in two
        // places. We need to set the COW flag to prevent mutations of one copy
        // affecting the other one.
        value = apfl_value_set_cow_flag(value);
    } else {
        stack_must_drop(ctx, -1);
    }
    return APFL_RESULT_OK;
}

static enum apfl_result
variable_new(apfl_ctx ctx, struct apfl_string *name)
{
    if (!apfl_scope_create_var(&ctx->gc, ctx->scope, name)) {
        return APFL_RESULT_ERR_FATAL;
    }
    return APFL_RESULT_OK;
}

static enum apfl_result
evaluate(apfl_ctx ctx, size_t *i, struct instruction_list *ilist)
{
    union instruction_or_arg arg;

    assert(*i < ilist->len);
    switch (ilist->instructions[(*i)++].instruction) {
    case INSN_NIL:
        return apfl_push_nil(ctx);
    case INSN_TRUE:
        return apfl_push_bool(ctx, true);
    case INSN_FALSE:
        return apfl_push_bool(ctx, false);
    case INSN_NUMBER:
        TRY(get_argument(i, ilist, &arg));
        return apfl_push_number(ctx, arg.number);
    case INSN_STRING:
        TRY(get_argument(i, ilist, &arg));
        return stack_push_or_fatal(ctx, (struct apfl_value) {
            .type = VALUE_STRING,
            .string = arg.string,
        });
    case INSN_LIST:
        TRY(get_argument(i, ilist, &arg));
        return apfl_list_create(ctx, arg.count);
    case INSN_LIST_APPEND:
        return apfl_list_append(ctx, -2, -1);
    case INSN_LIST_EXPAND_INTO:
        return apfl_list_append_list(ctx, -2, -1);
    case INSN_DICT:
        return apfl_dict_create(ctx);
    case INSN_DICT_APPEND_KVPAIR:
        return apfl_dict_set(ctx, -3, -2, -1);
    case INSN_GET_MEMBER:
        return apfl_get_member(ctx, -2, -1);
    case INSN_VAR_NEW:
        TRY(get_argument(i, ilist, &arg));
        return variable_new(ctx, arg.string);
    case INSN_VAR_GET:
        TRY(get_argument(i, ilist, &arg));
        return variable_get(ctx, arg.string);
    case INSN_VAR_SET:
        TRY(get_argument(i, ilist, &arg));
        return variable_set(ctx, arg.string, true);
    case INSN_NEXT_LINE:
        ctx->execution_line++;
        return APFL_RESULT_OK;
    case INSN_SET_LINE:
        TRY(get_argument(i, ilist, &arg));
        ctx->execution_line = arg.count;
        return APFL_RESULT_OK;
    }

    assert(false);
    return APFL_RESULT_ERR;
}

static enum apfl_result
evaluate_list(apfl_ctx ctx, struct instruction_list *ilist)
{
    ctx->execution_line = ilist->line;
    size_t i = 0;
    while (i < ilist->len) {
        TRY(evaluate(ctx, &i, ilist));
    }
    return APFL_RESULT_OK;
}

static enum apfl_result
eval_inner(apfl_ctx ctx, struct apfl_expr expr)
{
    struct instruction_list *ilist = apfl_instructions_new(&ctx->gc, expr.position.line);
    if (ilist == NULL) {
        return APFL_RESULT_ERR_FATAL;
    }

    if (!apfl_gc_tmproot_add(&ctx->gc, GC_OBJECT_FROM(ilist, GC_TYPE_INSTRUCTIONS))) {
        return APFL_RESULT_ERR_FATAL;
    }

    struct apfl_error error;
    if (!apfl_compile(&ctx->gc, expr, &error, ilist)) {
        return APFL_RESULT_ERR;
    }

    return evaluate_list(ctx, ilist);
}

static enum apfl_result
eval_expr(apfl_ctx ctx, struct apfl_expr expr)
{
    size_t tmproots = apfl_gc_tmproots_begin(&ctx->gc);
    enum apfl_result result = eval_inner(ctx, expr);
    apfl_gc_tmproots_restore(&ctx->gc, tmproots);

    return result;
}

bool
apfl_debug_print_val(apfl_ctx ctx, apfl_stackidx index, struct apfl_format_writer w)
{
    struct apfl_value value;
    if (!apfl_stack_pop(ctx, &value, index)) {
        FMT_TRY(apfl_format_put_string(w, "apfl_debug_print_val: Invalid stack index "));
        FMT_TRY(apfl_format_put_int(w, (int)index));
        FMT_TRY(apfl_format_put_string(w, "\n"));
        return true;
    }

    return apfl_value_print(value, w);
}

struct apfl_iterative_runner_data {
    apfl_ctx ctx;
    apfl_tokenizer_ptr tokenizer;
    apfl_parser_ptr parser;
    enum apfl_result result;
    bool end;
};

apfl_iterative_runner
apfl_iterative_runner_new(apfl_ctx ctx, struct apfl_source_reader reader)
{
    apfl_iterative_runner runner = ALLOC_OBJ(ctx->gc.allocator, struct apfl_iterative_runner_data);
    if (runner == NULL) {
        return NULL;
    }

    apfl_tokenizer_ptr tokenizer = apfl_tokenizer_new(ctx->gc.allocator, reader);
    if (tokenizer == NULL) {
        FREE_OBJ(ctx->gc.allocator, runner);
        return NULL;
    }

    apfl_parser_ptr parser = apfl_parser_new(ctx->gc.allocator, apfl_tokenizer_as_token_source(tokenizer));
    if (parser == NULL) {
        FREE_OBJ(ctx->gc.allocator, runner);
        apfl_tokenizer_destroy(tokenizer);
        return NULL;
    }

    *runner = (struct apfl_iterative_runner_data) {
        .ctx = ctx,
        .tokenizer = tokenizer,
        .parser = parser,
        .result = APFL_RESULT_OK,
        .end = false,
    };

    return runner;
}

static bool
move_string_onto_stack(apfl_ctx ctx, struct apfl_string string)
{
    struct apfl_value *value = apfl_stack_push_placeholder(ctx);
    if (value == NULL) {
        return false;
    }

    if ((value->string = apfl_string_move_into_new_gc_string(&ctx->gc, &string)) == NULL) {
        return false;
    }

    value->type = VALUE_STRING;

    return true;
}

static enum apfl_result
handle_parse_error(apfl_ctx ctx, struct apfl_error error)
{
    enum apfl_result result = APFL_RESULT_ERR;
    if (apfl_error_is_fatal_type(error.type)) {
        result = APFL_RESULT_ERR_FATAL;
    }

    const char *const_str = apfl_error_as_const_string(error);
    if (const_str != NULL) {
        return apfl_push_const_string(ctx, const_str)
            ? result
            : APFL_RESULT_ERR_FATAL;
    }

    struct apfl_string string;
    if (!apfl_error_as_string(error, ctx->gc.allocator, &string)) {
        // TODO: Maybe try to push something on the stack in case of error (also for other error cases below)?
        return APFL_RESULT_ERR_FATAL;
    }

    if (!move_string_onto_stack(ctx, string)) {
        return APFL_RESULT_ERR_FATAL;
    }

    return result;
}

bool
apfl_iterative_runner_next(apfl_iterative_runner runner)
{
    if (runner->end) {
        return false;
    }

    apfl_stack_clear(runner->ctx);

    switch (apfl_parser_next(runner->parser)) {
    case APFL_PARSE_OK:
        goto ok;
    case APFL_PARSE_ERROR:
        runner->result = handle_parse_error(runner->ctx, apfl_parser_get_error(runner->parser));
        return true;
    case APFL_PARSE_EOF:
        runner->end = true;
        return false;
    }

    assert(false);

ok:
    runner->result = eval_expr(runner->ctx, apfl_parser_get_expr(runner->parser));
    return true;
}

enum apfl_result
apfl_iterative_runner_get_result(apfl_iterative_runner runner)
{
    return runner->result;
}

void
apfl_iterative_runner_destroy(apfl_iterative_runner runner)
{
    if (runner == NULL) {
        return;
    }

    apfl_parser_destroy(runner->parser);
    apfl_tokenizer_destroy(runner->tokenizer);
    FREE_OBJ(runner->ctx->gc.allocator, runner);
}