#include <assert.h>
#include <stdbool.h>
#include <stdio.h>
#include <string.h>

#include "apfl.h"
#include "compile.h"
#include "bytecode.h"
#include "resizable.h"
#include "strings.h"

#define DEBUG_COMPILING 1

#if DEBUG_COMPILING
# define DEBUG_LOG(...) fprintf(stderr, __VA_ARGS__);
#else
# define DEBUG_LOG(...)
#endif

struct compiler {
    struct gc *gc;
    struct apfl_error error;
    size_t line;
};

#define TRY(b) do { if (!(b)) { return false; } } while(0)
#define MALLOC_FAIL_IF_NULL(compiler, x)                                 \
    do {                                                                 \
        if ((x) == NULL) {                                               \
            compiler->error = apfl_error_simple(APFL_ERR_MALLOC_FAILED); \
            return false;                                                \
        }                                                                \
    } while (0)

static bool compile_expr(struct compiler *, struct apfl_expr *, struct instruction_list *);

static bool
malloc_failure_on_false(struct compiler *compiler, bool b)
{
    if (!b) {
        compiler->error = apfl_error_simple(APFL_ERR_MALLOC_FAILED);
    }
    return b;
}

static bool
ilist_ensure_cap(struct compiler *compiler, struct instruction_list *ilist, size_t n)
{
    return malloc_failure_on_false(compiler, apfl_resizable_ensure_cap_for_more_elements(
        compiler->gc->allocator,
        sizeof(union instruction_or_arg),
        (void **)&ilist->instructions,
        ilist->len,
        &ilist->cap,
        n
    ));
}

static void
ilist_put_insn(struct instruction_list *ilist, enum instruction instruction)
{
    assert(ilist->cap > 0);
    assert(ilist->len < ilist->cap);

    ilist->instructions[ilist->len] = (union instruction_or_arg) {
        .instruction = instruction,
    };
    ilist->len++;

    DEBUG_LOG("put_insn: %s\n", apfl_instruction_to_string(instruction));
}

static void
ilist_put_number(struct instruction_list *ilist, apfl_number number)
{
    assert(ilist->cap > 0);
    assert(ilist->len < ilist->cap);

    ilist->instructions[ilist->len] = (union instruction_or_arg) {
        .number = number
    };
    ilist->len++;

    DEBUG_LOG("put_number: %lf\n", number);
}

static void
ilist_put_count(struct instruction_list *ilist, size_t count)
{
    assert(ilist->cap > 0);
    assert(ilist->len < ilist->cap);

    ilist->instructions[ilist->len] = (union instruction_or_arg) {
        .count = count
    };
    ilist->len++;

    DEBUG_LOG("put_count: %d\n", (int)count);
}

static void
ilist_put_string(struct instruction_list *ilist, struct apfl_string *string)
{
    assert(ilist->cap > 0);
    assert(ilist->len < ilist->cap);

    ilist->instructions[ilist->len] = (union instruction_or_arg) {
        .string = string
    };
    ilist->len++;

    DEBUG_LOG("put_string: " APFL_STR_FMT "\n", APFL_STR_FMT_ARGS(*string));
}

static bool
string_move_into_new(struct compiler *compiler, struct apfl_string **out, struct apfl_string *in)
{
    struct apfl_string *str = apfl_string_move_into_new_gc_string(compiler->gc, in);
    if (str == NULL) {
        compiler->error = apfl_error_simple(APFL_ERR_MALLOC_FAILED);
        return false;
    }

    *out = str;
    return true;
}

static bool
compile_constant(struct compiler *compiler, struct apfl_expr_const *constant, struct instruction_list *ilist)
{
    switch (constant->type) {
    case APFL_EXPR_CONST_NIL:
        TRY(ilist_ensure_cap(compiler, ilist, 1));
        ilist_put_insn(ilist, INSN_NIL);
        return true;
    case APFL_EXPR_CONST_BOOLEAN:
        TRY(ilist_ensure_cap(compiler, ilist, 1));
        ilist_put_insn(ilist, constant->boolean ? INSN_TRUE : INSN_FALSE);
        return true;
    case APFL_EXPR_CONST_NUMBER:
        TRY(ilist_ensure_cap(compiler, ilist, 2));
        ilist_put_insn(ilist, INSN_NUMBER);
        ilist_put_number(ilist, constant->number);
        return true;
    case APFL_EXPR_CONST_STRING:
        TRY(ilist_ensure_cap(compiler, ilist, 2));

        struct apfl_string *str;
        TRY(string_move_into_new(compiler, &str, &constant->string));

        ilist_put_insn(ilist, INSN_STRING);
        ilist_put_string(ilist, str);
        return true;
    }

    assert(false);
    return false;
}

static bool
compile_list(struct compiler *compiler, struct apfl_expr_list *list, struct instruction_list *ilist)
{
    TRY(ilist_ensure_cap(compiler, ilist, 2));
    ilist_put_insn(ilist, INSN_LIST);
    ilist_put_count(ilist, list->len);

    for (size_t i = 0; i < list->len; i++) {
        struct apfl_expr_list_item *item = &list->items[i];

        TRY(compile_expr(compiler, item->expr, ilist));
        TRY(ilist_ensure_cap(compiler, ilist, 1));
        ilist_put_insn(ilist, item->expand ? INSN_LIST_EXPAND_INTO : INSN_LIST_APPEND);
    }

    return true;
}

static bool
compile_dict(struct compiler *compiler, struct apfl_expr_dict *dict, struct instruction_list *ilist)
{
    TRY(ilist_ensure_cap(compiler, ilist, 1));
    ilist_put_insn(ilist, INSN_DICT);

    for (size_t i = 0; i < dict->len; i++) {
        struct apfl_expr_dict_pair *pair = &dict->items[i];

        TRY(compile_expr(compiler, pair->k, ilist));
        TRY(compile_expr(compiler, pair->v, ilist));
        TRY(ilist_ensure_cap(compiler, ilist, 1));
        ilist_put_insn(ilist, INSN_DICT_APPEND_KVPAIR);
    }

    return true;
}

static bool
compile_dot(struct compiler *compiler, struct apfl_expr_dot *dot, struct instruction_list *ilist)
{
    TRY(compile_expr(compiler, dot->lhs, ilist));

    TRY(ilist_ensure_cap(compiler, ilist, 3));

    struct apfl_string *str;
    TRY(string_move_into_new(compiler, &str, &dot->rhs));

    ilist_put_insn(ilist, INSN_STRING);
    ilist_put_string(ilist, str);
    ilist_put_insn(ilist, INSN_GET_MEMBER);

    return true;
}

static bool
compile_at(struct compiler *compiler, struct apfl_expr_at *at, struct instruction_list *ilist)
{
    TRY(compile_expr(compiler, at->lhs, ilist));
    TRY(compile_expr(compiler, at->rhs, ilist));

    TRY(ilist_ensure_cap(compiler, ilist, 1));
    ilist_put_insn(ilist, INSN_GET_MEMBER);

    return true;
}

static bool
compile_var(struct compiler *compiler, struct apfl_string *var, struct instruction_list *ilist)
{
    TRY(ilist_ensure_cap(compiler, ilist, 2));

    struct apfl_string *str;
    TRY(string_move_into_new(compiler, &str, var));

    ilist_put_insn(ilist, INSN_VAR_GET);
    ilist_put_string(ilist, str);

    return true;
}

static bool
compile_simple_assignment(
    struct compiler *compiler,
    struct apfl_string *var,
    struct apfl_expr *rhs,
    struct instruction_list *ilist
) {
    TRY(ilist_ensure_cap(compiler, ilist, 2));

    struct apfl_string *str;
    TRY(string_move_into_new(compiler, &str, var));

    ilist_put_insn(ilist, INSN_VAR_NEW);
    ilist_put_string(ilist, str);

    TRY(compile_expr(compiler, rhs, ilist));

    TRY(ilist_ensure_cap(compiler, ilist, 2));
    ilist_put_insn(ilist, INSN_VAR_SET);
    ilist_put_string(ilist, str);

    return true;
}

static bool
compile_assignment(
    struct compiler *compiler,
    struct apfl_expr_assignment *assignment,
    struct instruction_list *ilist
) {
    if (
        assignment->lhs.type == APFL_EXPR_ASSIGNABLE_VAR_OR_MEMBER
        && assignment->lhs.var_or_member.type == APFL_EXPR_ASSIGNABLE_VAR_OR_MEMBER_VAR
    ) {
        return compile_simple_assignment(
            compiler,
            &assignment->lhs.var_or_member.var,
            assignment->rhs,
            ilist
        );
    }

    // TODO: Implement other assignables
    compiler->error = apfl_error_simple(APFL_ERR_NOT_IMPLEMENTED);
    return false;
}

static bool
compile_expr(struct compiler *compiler, struct apfl_expr *expr, struct instruction_list *ilist)
{
    size_t new_line = (size_t)expr->position.line;
    if (new_line != compiler->line) {
        if (new_line == compiler->line + 1) {
            TRY(ilist_ensure_cap(compiler, ilist, 1));
            ilist_put_insn(ilist, INSN_NEXT_LINE);
        } else {
            TRY(ilist_ensure_cap(compiler, ilist, 2));
            ilist_put_insn(ilist, INSN_SET_LINE);
            ilist_put_count(ilist, new_line);
        }

        compiler->line = new_line;
    }

    switch (expr->type) {
    case APFL_EXPR_CALL:
    case APFL_EXPR_SIMPLE_FUNC:
    case APFL_EXPR_COMPLEX_FUNC:
        compiler->error = apfl_error_simple(APFL_ERR_NOT_IMPLEMENTED);
        return false;
    case APFL_EXPR_CONSTANT:
        return compile_constant(compiler, &expr->constant, ilist);
    case APFL_EXPR_BLANK:
        TRY(ilist_ensure_cap(compiler, ilist, 1));
        ilist_put_insn(ilist, INSN_NIL);
        return true;
    case APFL_EXPR_LIST:
        return compile_list(compiler, &expr->list, ilist);
    case APFL_EXPR_DICT:
        return compile_dict(compiler, &expr->dict, ilist);
    case APFL_EXPR_DOT:
        return compile_dot(compiler, &expr->dot, ilist);
    case APFL_EXPR_AT:
        return compile_at(compiler, &expr->at, ilist);
    case APFL_EXPR_VAR:
        return compile_var(compiler, &expr->var, ilist);
    case APFL_EXPR_ASSIGNMENT:
        return compile_assignment(compiler, &expr->assignment, ilist);
    }

    assert(false);
    return false;
}

bool
apfl_compile(struct gc *gc, struct apfl_expr expr, struct apfl_error *error_out, struct instruction_list *out)
{
    struct compiler compiler = {
        .gc = gc,
        .line = out->line,
    };

    bool ok = compile_expr(&compiler, &expr, out);
    apfl_expr_deinit(gc->allocator, &expr);

    if (!ok) {
        *error_out = compiler.error;
    }
    return ok;
}