#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"

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

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

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

static void
must_get_argument(apfl_ctx ctx, size_t *i, struct instruction_list *ilist, union instruction_or_arg *arg)
{
    if (!get_argument(i, ilist, arg)) {
        apfl_raise_const_error(ctx, APFL_RESULT_ERR, apfl_messages.corrupted_bytecode);
    }
}

static void
variable_get(apfl_ctx ctx, struct apfl_string *name)
{
    struct apfl_value value;
    if (!apfl_scope_get(ctx->scope, name, &value)) {
        // TODO: Include variable name in error
        apfl_raise_const_error(ctx, APFL_RESULT_ERR, apfl_messages.variable_doesnt_exist);
    }

    apfl_stack_must_push(ctx, value);
}

static void
variable_set(apfl_ctx ctx, struct apfl_string *name, bool keep_on_stack)
{
    struct apfl_value value;
    if (!apfl_stack_get(ctx, &value, -1)) {
        apfl_raise_invalid_stackidx(ctx);
    }
    if (!apfl_scope_set(&ctx->gc, ctx->scope, name, value)) {
        apfl_raise_alloc_error(ctx);
    }
    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);
    }
}

static void
variable_new(apfl_ctx ctx, struct apfl_string *name)
{
    if (!apfl_scope_create_var(&ctx->gc, ctx->scope, name)) {
        apfl_raise_alloc_error(ctx);
    }
}

static void
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:
        apfl_push_nil(ctx);
        return;
    case INSN_TRUE:
        apfl_push_bool(ctx, true);
        return;
    case INSN_FALSE:
        apfl_push_bool(ctx, false);
        return;
    case INSN_NUMBER:
        must_get_argument(ctx, i, ilist, &arg);
        apfl_push_number(ctx, arg.number);
        return;
    case INSN_STRING:
        must_get_argument(ctx, i, ilist, &arg);
        apfl_stack_must_push(ctx, (struct apfl_value) {
            .type = VALUE_STRING,
            .string = arg.string,
        });
        return;
    case INSN_LIST:
        must_get_argument(ctx, i, ilist, &arg);
        apfl_list_create(ctx, arg.count);
        return;
    case INSN_LIST_APPEND:
        apfl_list_append(ctx, -2, -1);
        return;
    case INSN_LIST_EXPAND_INTO:
        apfl_list_append_list(ctx, -2, -1);
        return;
    case INSN_DICT:
        apfl_dict_create(ctx);
        return;
    case INSN_DICT_APPEND_KVPAIR:
        apfl_dict_set(ctx, -3, -2, -1);
        return;
    case INSN_GET_MEMBER:
        apfl_get_member(ctx, -2, -1);
        return;
    case INSN_VAR_NEW:
        must_get_argument(ctx, i, ilist, &arg);
        variable_new(ctx, arg.string);
        return;
    case INSN_VAR_GET:
        must_get_argument(ctx, i, ilist, &arg);
        variable_get(ctx, arg.string);
        return;
    case INSN_VAR_SET:
        must_get_argument(ctx, i, ilist, &arg);
        variable_set(ctx, arg.string, true);
        return;
    case INSN_NEXT_LINE:
        ctx->execution_line++;
        return;
    case INSN_SET_LINE:
        must_get_argument(ctx, i, ilist, &arg);
        ctx->execution_line = arg.count;
        return;
    }

    assert(false);
}

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

static void
eval_expr_inner(apfl_ctx ctx, struct apfl_expr expr)
{
    struct instruction_list *ilist = apfl_instructions_new(&ctx->gc, expr.position.line);
    if (ilist == NULL) {
        apfl_raise_alloc_error(ctx);
    }

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

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

    evaluate_list(ctx, ilist);
}

static void
eval_expr(apfl_ctx ctx, struct apfl_expr expr)
{
    size_t tmproots = apfl_gc_tmproots_begin(&ctx->gc);
    eval_expr_inner(ctx, expr);
    apfl_gc_tmproots_restore(&ctx->gc, tmproots);
}

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 with_error_on_stack;
    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 void
handle_parse_error(apfl_ctx ctx, struct apfl_error error)
{
    enum apfl_result errtype = APFL_RESULT_ERR;
    if (apfl_error_is_fatal_type(error.type)) {
        errtype = APFL_RESULT_ERR_FATAL;
    }

    const char *const_str = apfl_error_as_const_string(error);
    if (const_str != NULL) {
        apfl_raise_const_error(ctx, errtype, const_str);
    }

    struct apfl_string string;
    if (!apfl_error_as_string(error, ctx->gc.allocator, &string)) {
        apfl_raise_alloc_error(ctx);
    }

    if (!apfl_move_string_onto_stack(ctx, string)) {
        apfl_raise_alloc_error(ctx);
    }

    apfl_raise_error_with_type(ctx, -1, errtype);
}

static void
iterative_runner_next_protected(apfl_ctx ctx, void *opaque)
{
    apfl_iterative_runner runner = opaque;

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

    assert(false);
}

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

    apfl_stack_clear(runner->ctx);

    runner->with_error_on_stack = false;
    runner->result = apfl_protected(
        runner->ctx,
        iterative_runner_next_protected,
        runner,
        &runner->with_error_on_stack
    );

    return !runner->end;
}

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

bool
apfl_iterative_runner_has_error_on_stack(apfl_iterative_runner runner)
{
    return runner->with_error_on_stack;
}

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);
}