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

#include "apfl.h"

#include "alloc.h"
#include "bytecode.h"
#include "context.h"
#include "globals.h"
#include "scope.h"

#define TRY_FORMAT(ctx, x)                                         \
    do {                                                           \
        if (!(x)) {                                                \
            apfl_raise_const_error((ctx), apfl_messages.io_error); \
        }                                                          \
    } while (0)

static void
call_then_or_else(apfl_ctx ctx, size_t arg)
{
    apfl_get_list_member_by_index(ctx, 0, arg);
    apfl_list_create(ctx, 0);
    apfl_call(ctx, -2, -1);
}

static void
impl_if(apfl_ctx ctx)
{
    size_t len = apfl_len(ctx, 0);

    if (len < 2) {
        apfl_raise_const_error(ctx, "if needs at least 2 arguments");
    }

    size_t i = 0;
    while (i+1 < len) {
        size_t cond = i++;
        size_t then = i++;

        apfl_get_list_member_by_index(ctx, 0, cond);
        if (apfl_get_type(ctx, -1) == APFL_VALUE_FUNC) {
            apfl_list_create(ctx, 0);
            apfl_call(ctx, -2, -1);
        }

        if (apfl_is_truthy(ctx, -1)) {
            apfl_get_list_member_by_index(ctx, 0, then);
            apfl_list_create(ctx, 0);
            apfl_call(ctx, -2, -1);
            return;
        }
    }

    if (i == len) {
        // Empty else. Return nil
        apfl_push_nil(ctx);
    } else {
        assert(i+1 == len /*exactly one argument remains as final else*/);

        call_then_or_else(ctx, i);
    }
}

static void
impl_eq(apfl_ctx ctx)
{
    size_t len = apfl_len(ctx, 0);
    if (len < 2) {
        apfl_raise_const_error(ctx, "== needs at least 2 arguments");
    }

    for (size_t i = 0; i < len-1; i++) {
        apfl_get_list_member_by_index(ctx, 0, i);
        apfl_get_list_member_by_index(ctx, 0, i+1);
        if (!apfl_eq(ctx, -2, -1)) {
            apfl_push_bool(ctx, false);
            return;
        }
    }

    apfl_push_bool(ctx, true);
}

#define IMPLEMENT_COMPARISON(impl_name, name, cmp)                                           \
    static void                                                                              \
    impl_name(apfl_ctx ctx)                                                                  \
    {                                                                                        \
        size_t len = apfl_len(ctx, 0);                                                       \
        if (len < 2) {                                                                       \
            apfl_raise_const_error(ctx, name " needs at least 2 arguments"); \
        }                                                                                    \
                                                                                             \
        for (size_t i = 0; i < len-1; i++) {                                                 \
            apfl_get_list_member_by_index(ctx, 0, i);                                        \
            apfl_get_list_member_by_index(ctx, 0, i+1);                                      \
            if (!(apfl_cmp(ctx, -2, -1) cmp 0)) {                                            \
                apfl_push_bool(ctx, false);                                                  \
                return;                                                                      \
            }                                                                                \
        }                                                                                    \
                                                                                             \
        apfl_push_bool(ctx, true);                                                           \
    }                                                                                        \

IMPLEMENT_COMPARISON(impl_gt, ">", >)
IMPLEMENT_COMPARISON(impl_lt, "<", <)
IMPLEMENT_COMPARISON(impl_ge, ">=", >=)
IMPLEMENT_COMPARISON(impl_le, "<=", <=)

#define IMPL_MATH_OP(name, op)                       \
    static apfl_number                               \
    name(apfl_ctx ctx, apfl_number a, apfl_number b) \
    {                                                \
        (void)ctx;                                   \
        return a op b;                               \
    }

#define IMPL_MATH_OP_FUNC(impl, name, single, op)                                            \
    static void                                                                              \
    impl(apfl_ctx ctx)                                                                       \
    {                                                                                        \
        size_t len = apfl_len(ctx, 0);                                                       \
        if (len < 1) {                                                                       \
            apfl_raise_const_error(ctx, name " needs at least 1 argument"); \
        }                                                                                    \
                                                                                             \
        apfl_get_list_member_by_index(ctx, 0, 0);                                            \
        apfl_number num = apfl_get_number(ctx, -1);                                          \
        if (len == 1) {                                                                      \
            apfl_push_number(ctx, single(num));                                              \
            return;                                                                          \
        }                                                                                    \
                                                                                             \
        for (size_t i = 1; i < len; i++) {                                                   \
            apfl_get_list_member_by_index(ctx, 0, i);                                        \
            num = op(ctx, num, apfl_get_number(ctx, -1));                                    \
        }                                                                                    \
                                                                                             \
        apfl_push_number(ctx, num);                                                          \
    }

static apfl_number
single_identity(apfl_number number)
{
    return number;
}

static apfl_number
single_negate(apfl_number number)
{
    return -number;
}

IMPL_MATH_OP(op_plus, +)
IMPL_MATH_OP(op_minus, -)
IMPL_MATH_OP(op_mult, *)
IMPL_MATH_OP_FUNC(impl_plus,  "+", single_identity, op_plus)
IMPL_MATH_OP_FUNC(impl_minus, "-", single_negate,   op_minus)
IMPL_MATH_OP_FUNC(impl_mult,  "*", single_identity, op_mult)

static apfl_number
op_div(apfl_ctx ctx, apfl_number a, apfl_number b)
{
    if (b == 0.0 || b == -0.0) { // apfl_number is a double, so there are technically two zeroes :/
        apfl_raise_const_error(ctx, "Division by zero");
    }
    return a / b;
}

IMPL_MATH_OP_FUNC(impl_div,   "/", single_identity, op_div)

static void
impl_concat(apfl_ctx ctx)
{
    apfl_push_const_string(ctx, "");
    apfl_join_strings(ctx, -1, -2);
}

static void
impl_join(apfl_ctx ctx)
{
    size_t len = apfl_len(ctx, 0);
    if (len != 2) {
        apfl_raise_const_error(ctx, "join expects exactly 2 arguments");
    }

    apfl_get_list_member_by_index(ctx, 0, 0);
    apfl_get_list_member_by_index(ctx, 0, 1);
    apfl_drop(ctx, 0);
    apfl_join_strings(ctx, -2, -1);
}

static void
print(apfl_ctx ctx)
{
    struct apfl_io_writer w = apfl_get_output_writer(ctx);

    size_t len = apfl_len(ctx, 0);
    for (size_t i = 0; i < len; i++) {
        if (i > 0) {
            TRY_FORMAT(ctx, apfl_io_write_string(w, " "));
        }

        apfl_get_list_member_by_index(ctx, 0, i);
        apfl_tostring(ctx, -1);
        struct apfl_string_view sv = apfl_get_string(ctx, -1);
        TRY_FORMAT(ctx, apfl_io_write_string(w, sv));
        apfl_drop(ctx, -1);
    }

    if (len > 0) {
        TRY_FORMAT(ctx, apfl_io_write_string(w, "\n"));
    }

    apfl_push_nil(ctx);
}

static void
dump(apfl_ctx ctx)
{
    struct apfl_io_writer w = apfl_get_output_writer(ctx);
    apfl_get_list_member_by_index(ctx, 0, 0);
    apfl_drop(ctx, 0);
    TRY_FORMAT(ctx, apfl_debug_print_val(ctx, -1, w));
}

static void
disasm(apfl_ctx ctx)
{
    struct apfl_io_writer w = apfl_get_output_writer(ctx);
    apfl_get_list_member_by_index(ctx, 0, 0);
    apfl_drop(ctx, 0);
    struct apfl_value value = apfl_stack_must_get(ctx, -1);
    if (value.type == VALUE_CFUNC) {
        apfl_raise_const_error(ctx, "disasm needs a apfl function, got a native function instead");
    } else if (value.type != VALUE_FUNC) {
        apfl_raise_errorfmt(ctx, "disasm needs a apfl function, got value of type {value:type} instead", value);
    }

    struct function *func = value.func;

    for (size_t i = 0; i < func->subfunctions_len; i++) {
        struct subfunction *subfunction = &func->subfunctions[i];

        TRY_FORMAT(ctx, apfl_io_write_string(w, "Subfunction #"));
        TRY_FORMAT(ctx, apfl_format_put_number(w, (int)i));
        TRY_FORMAT(ctx, apfl_io_write_string(w, "\n"));
        if (subfunction->matcher == NULL) {
            TRY_FORMAT(ctx, apfl_format_put_indent(w, 1));
            TRY_FORMAT(ctx, apfl_io_write_string(w, "No matcher (matches everything)\n"));
        } else {
            TRY_FORMAT(ctx, apfl_format_put_indent(w, 1));
            TRY_FORMAT(ctx, apfl_io_write_string(w, "Matcher\n"));
            TRY_FORMAT(ctx, apfl_bytecode_dump_matcher(2, w, subfunction->matcher->instructions));
        }
        TRY_FORMAT(ctx, apfl_format_put_indent(w, 1));
        TRY_FORMAT(ctx, apfl_io_write_string(w, "Instructions\n"));
        TRY_FORMAT(ctx, apfl_bytecode_dump(2, w, subfunction->body));
    }
}

static void
require_exactly_one_arg(apfl_ctx ctx, const char *error)
{
    size_t len = apfl_len(ctx, 0);
    if (len != 1) {
        apfl_raise_const_error(ctx, error);
    }

    apfl_get_list_member_by_index(ctx, 0, 0);
    apfl_drop(ctx, 0);
}

#define ONE_ARG(ctx, name) require_exactly_one_arg((ctx), name " needs exactly 1 argument")

static void
tostring(apfl_ctx ctx)
{
    ONE_ARG(ctx, "tostring");
    apfl_tostring(ctx, -1);
}

static void
not(apfl_ctx ctx)
{
    ONE_ARG(ctx, "not");
    bool b = apfl_is_truthy(ctx, -1);
    apfl_push_bool(ctx, !b);
}

static void
len(apfl_ctx ctx)
{
    ONE_ARG(ctx, "len");
    size_t l = apfl_len(ctx, -1);
    apfl_push_number(ctx, (apfl_number)l);
}

static void
type(apfl_ctx ctx)
{
    ONE_ARG(ctx, "type");
    apfl_push_const_string(ctx, apfl_type_name(apfl_get_type(ctx, -1)));
}

static void
impl_while(apfl_ctx ctx)
{
    size_t argc = apfl_len(ctx, 0);
    if (argc != 2) {
        apfl_raise_const_error(ctx, "while needs 2 functions as arguments");
    }

    apfl_get_list_member_by_index(ctx, 0, 0);
    if (apfl_get_type(ctx, -1) != APFL_VALUE_FUNC) {
        apfl_raise_const_error(ctx, "while needs 2 functions as arguments");
    }
    apfl_get_list_member_by_index(ctx, 0, 1);
    if (apfl_get_type(ctx, -1) != APFL_VALUE_FUNC) {
        apfl_raise_const_error(ctx, "while needs 2 functions as arguments");
    }

    apfl_drop(ctx, 0);

    apfl_push_nil(ctx); // Return value in case of no iteration

    for (;;) {
        apfl_copy(ctx, 0);
        apfl_list_create(ctx, 0);
        apfl_call(ctx, -2, -1);
        if (!apfl_is_truthy(ctx, -1)) {
            break;
        }

        apfl_drop(ctx, -1);
        apfl_copy(ctx, 1);
        apfl_list_create(ctx, 0);
        apfl_call(ctx, -2, -1);
    }
}

static void
impl_gc(apfl_ctx ctx)
{
    ONE_ARG(ctx, "gc");
    apfl_tostring(ctx, -1);
    struct apfl_string_view sv = apfl_get_string(ctx, -1);

    if (apfl_string_eq(sv, "dump")) {
        struct apfl_io_writer w = apfl_get_output_writer(ctx);
        TRY_FORMAT(ctx, apfl_gc_debug_dump_graph(&ctx->gc, w));
    } else if (apfl_string_eq(sv, "collect")) {
        apfl_gc_full(&ctx->gc);
    } else {
        apfl_raise_const_error(ctx, "Unknown gc command");
    }

    apfl_drop(ctx, -1);
}

static void
impl_backtrace(apfl_ctx ctx)
{
    apfl_drop(ctx, -1);

    struct apfl_io_writer w = apfl_get_output_writer(ctx);

    size_t depth = apfl_call_stack_depth(ctx);
    for (size_t i = 1; i < depth; i++) {
        TRY_FORMAT(ctx, apfl_io_write_string(w, "#"));
        TRY_FORMAT(ctx, apfl_format_put_int(w, (int)i));
        TRY_FORMAT(ctx, apfl_io_write_string(w, ": "));
        TRY_FORMAT(ctx, apfl_call_stack_entry_info_format(
            w,
            apfl_call_stack_inspect(ctx, i)
        ));
        TRY_FORMAT(ctx, apfl_io_write_string(w, "\n"));
    }
}

static void
closefile(FILE **f)
{
    if (*f != NULL) {
        fclose(*f);
        *f = NULL;
    }
}

static void
file_onbeforecollect(void *opaque)
{
    FILE **f = opaque;
    closefile(f);
}

static struct apfl_native_object_type file_object = {
    .size = sizeof(FILE *),
    .onbeforecollect = file_onbeforecollect,
};

static const char *
getcstring(apfl_ctx ctx, apfl_stackidx index)
{
    apfl_move_to_top_of_stack(ctx, index);
    apfl_push_string_view_copy(ctx, (struct apfl_string_view) { .bytes = (unsigned char [1]) {'\0'}, .len = 1, });
    apfl_concat_strings(ctx, -2, -1);
    return (const char *)apfl_get_string(ctx, -1).bytes;
}

static void
raise_errno(apfl_ctx ctx)
{
    char *err = strerror(errno);
    apfl_push_string_view_copy(ctx, apfl_string_view_from(err));
    apfl_raise_error(ctx, -1);
}

static void
impl_fopen(apfl_ctx ctx)
{
    size_t argc = apfl_len(ctx, 0);
    if (argc == 0) {
        apfl_raise_const_error(ctx, "fopen needs at least one argument");
    }

    apfl_get_list_member_by_index(ctx, 0, 0);
    const char *filename = getcstring(ctx, -1);

    if (argc > 1) {
        apfl_get_list_member_by_index(ctx, 0, 1);
    } else {
        apfl_push_const_string(ctx, "rb");
    }
    const char *mode = getcstring(ctx, -1);

    FILE **fh = apfl_push_native_object(ctx, &file_object);

    *fh = fopen(filename, mode);
    if (*fh == NULL) {
        raise_errno(ctx);
    }
}

static void
impl_fread(apfl_ctx ctx)
{
    size_t argc = apfl_len(ctx, 0);
    if (argc != 2) {
        apfl_raise_const_error(ctx, "fread needs exactly 2 arguments");
    }

    apfl_get_list_member_by_index(ctx, 0, 0);
    FILE **fh = apfl_get_native_object(ctx, &file_object, -1);
    if (*fh == NULL) {
        apfl_raise_const_error(ctx, "File already closed");
    }

    apfl_get_list_member_by_index(ctx, 0, 1);
    apfl_number presize = apfl_get_number(ctx, -1);
    if (presize < 0) {
        apfl_raise_const_error(ctx, "Can not read a negative amount of bytes from file");
    }

    size_t size = (size_t)presize;

    unsigned char *buf = ALLOC_BYTES(ctx->gc.allocator, size);
    if (buf == NULL) {
        apfl_raise_alloc_error(ctx);
    }

    size_t actual = fread(buf, 1, size, *fh);

    if (!apfl_move_string_onto_stack(ctx, (struct apfl_string) {
        .bytes = buf,
        .len = actual,
        .cap = size,
    })) {
        FREE_BYTES(ctx->gc.allocator, buf, size);
    }
}

static void
impl_fwrite(apfl_ctx ctx)
{
    size_t argc = apfl_len(ctx, 0);
    if (argc != 2) {
        apfl_raise_const_error(ctx, "fread needs exactly 2 arguments");
    }

    apfl_get_list_member_by_index(ctx, 0, 0);
    FILE **fh = apfl_get_native_object(ctx, &file_object, -1);
    if (*fh == NULL) {
        apfl_raise_const_error(ctx, "File already closed");
    }

    apfl_get_list_member_by_index(ctx, 0, 1);
    struct apfl_string_view sv = apfl_get_string(ctx, -1);
    if (fwrite(sv.bytes, 1, sv.len, *fh) != sv.len) {
        raise_errno(ctx);
    }

    apfl_push_nil(ctx);
}

static void
impl_fclose(apfl_ctx ctx)
{
    ONE_ARG(ctx, "fclose");
    FILE **fh = apfl_get_native_object(ctx, &file_object, -1);
    closefile(fh);
    apfl_drop(ctx, -1);
}

static void
loadfile(apfl_ctx ctx)
{
    ONE_ARG(ctx, "loadfile");
    apfl_copy(ctx, -1);
    const char *filename = getcstring(ctx, -1);

    FILE **fh = apfl_push_native_object(ctx, &file_object);
    *fh = fopen(filename, "rb");
    if (*fh == NULL) {
        raise_errno(ctx);
    }
    apfl_drop(ctx, -2); // drop cstring

    struct apfl_io_reader r = apfl_io_file_reader(*fh);
    apfl_load(ctx, apfl_io_reader_as_source_reader(&r), -2);
    closefile(fh);
    apfl_drop(ctx, -2);
}

static void
loadstring(apfl_ctx ctx)
{
    ONE_ARG(ctx, "loadstring");
    apfl_tostring(ctx, -1);
    apfl_push_const_string(ctx, "(loadstring)");

    struct apfl_io_string_reader_data reader_data = apfl_io_string_reader_create(apfl_get_string(ctx, -2));
    struct apfl_io_reader r = apfl_io_string_reader(&reader_data);
    apfl_load(ctx, apfl_io_reader_as_source_reader(&r), -1);
    apfl_drop(ctx, -2);
}

static void
serialize_bytecode(apfl_ctx ctx)
{
    apfl_get_list_member_by_index(ctx, 0, 0);
    struct apfl_value value = apfl_stack_must_get(ctx, -1);
    if (value.type == VALUE_CFUNC) {
        apfl_raise_const_error(ctx, "-serialize-bytecode needs a apfl function, got a native function instead");
    } else if (value.type != VALUE_FUNC) {
        apfl_raise_errorfmt(ctx, "-serialize-bytecode needs a apfl function, got value of type {value:type} instead", value);
    }

    apfl_get_list_member_by_index(ctx, 0, 1);
    FILE **fh = apfl_get_native_object(ctx, &file_object, -1);

    struct apfl_io_writer w = apfl_io_file_writer(*fh);

    if (!apfl_bytecode_serialize(ctx->gc.allocator, w, value.func->subfunctions[0].body)) {
        apfl_raise_const_error(ctx, "Could not serialize function");
    }

    apfl_push_nil(ctx);
}

static void
unserialize_bytecode(apfl_ctx ctx)
{
    apfl_get_list_member_by_index(ctx, 0, 0);
    apfl_drop(ctx, -2);
    FILE **fh = apfl_get_native_object(ctx, &file_object, -1);

    struct apfl_io_reader r = apfl_io_file_reader(*fh);

    apfl_load_bytecode(ctx, r);
}

static const struct global_def globals[] = {
    {"if", impl_if},
    {"==", impl_eq},
    {">", impl_gt},
    {"<", impl_lt},
    {">=", impl_ge},
    {"<=", impl_le},
    {"+", impl_plus},
    {"-", impl_minus},
    {"*", impl_mult},
    {"/", impl_div},
    {"&", impl_concat},
    {"join", impl_join},
    {"print", print},
    {"dump", dump},
    {"disasm", disasm},
    {"tostring", tostring},
    {"not", not},
    {"len", len},
    {"type", type},
    {"while", impl_while},
    {"gc", impl_gc},
    {"backtrace", impl_backtrace},
    {"fopen", impl_fopen},
    {"fread", impl_fread},
    {"fwrite", impl_fwrite},
    {"fclose", impl_fclose},
    {"loadfile", loadfile},
    {"loadstring", loadstring},
    {"-serialize-bytecode", serialize_bytecode},
    {"-unserialize-bytecode", unserialize_bytecode},
    {NULL, NULL},
};

const struct global_def *
apfl_globals(void)
{
    return globals;
}