#include <assert.h>

#include "apfl.h"

#include "alloc.h"
#include "context.h"
#include "format.h"
#include "resizable.h"
#include "hashmap.h"
#include "value.h"

#define TRY FMT_TRY

size_t
apfl_list_len(struct list_header *list)
{
    return list->len;
}

static bool
dict_keys_eq(void *opaque, const void *a, const void *b)
{
    (void)opaque;
    return apfl_value_eq(*((struct apfl_value *) a), *((struct apfl_value *) b));
}

static apfl_hash
dict_calc_hash(void *opaque, const void *key)
{
    (void)opaque;
    return apfl_value_hash(*(const struct apfl_value *)key);
}

static bool
dict_init_hashmap(struct apfl_allocator allocator, struct apfl_hashmap *map)
{
    bool ok = apfl_hashmap_init(
        map,
        allocator,
        (struct apfl_hashmap_callbacks) {
            .keys_eq       = dict_keys_eq,
            .calc_hash     = dict_calc_hash,
        },
        sizeof(struct apfl_value),
        sizeof(struct apfl_value)
    );

    return ok;
}

struct dict_header *
apfl_dict_new(struct gc *gc)
{
    struct apfl_hashmap map;

    if (!dict_init_hashmap(gc->allocator, &map)) {
        return NULL;
    }

    struct dict_header *dict = apfl_gc_new_dict(gc);
    if (dict == NULL) {
        apfl_hashmap_deinit(&map);
        return NULL;
    }
    *dict = (struct dict_header) {
        .map = map,
        .copy_on_write = false,
    };
    return dict;
}

void
apfl_dict_deinit(struct dict_header *header)
{
    apfl_hashmap_deinit(&header->map);
}

static struct apfl_string_view
as_string_view(struct apfl_value value)
{
    switch (value.type) {
    case VALUE_STRING:
        return apfl_string_view_from(*value.string);
    case VALUE_CONST_STRING:
        return value.const_string;
    default:
        return (struct apfl_string_view) { .bytes = NULL, .len = 0 };
    }
}

static bool
format(unsigned indent, struct apfl_format_writer w, struct apfl_value value, bool skip_first_indent)
{
    TRY(apfl_format_put_indent(w, skip_first_indent ? 0 : indent));

    switch (value.type) {
    case VALUE_NIL:
        TRY(apfl_format_put_string(w, "nil"));
        return true;
    case VALUE_BOOLEAN:
        TRY(apfl_format_put_string(w, value.boolean ? "true" : "false"));
        return true;
    case VALUE_NUMBER:
        TRY(apfl_format_put_number(w, value.number));
        return true;
    case VALUE_STRING:
    case VALUE_CONST_STRING:
        TRY(apfl_format_put_string(w, "\""));
        TRY(apfl_format_put_string(w, as_string_view(value)));
        TRY(apfl_format_put_string(w, "\""));
        return true;
    case VALUE_LIST:
        if (value.list->len == 0) {
            return apfl_format_put_string(w, "[]");
        }
        TRY(apfl_format_put_string(w, "[\n"));
        for (size_t i = 0; i < value.list->len; i++) {
            TRY(format(indent+1, w, value.list->items[i], false));
            TRY(apfl_format_put_string(w, "\n"));
        }
        TRY(apfl_format_put_indent(w, indent));
        TRY(apfl_format_put_string(w, "]"));
        return true;
    case VALUE_DICT:
        if (apfl_hashmap_count(value.dict->map) == 0) {
            return apfl_format_put_string(w, "[->]");
        }

        TRY(apfl_format_put_string(w, "[\n"));

        HASHMAP_EACH(&value.dict->map, cur) {
            struct apfl_value *k = apfl_hashmap_cursor_peek_key(cur);
            assert(k != NULL);
            struct apfl_value *v = apfl_hashmap_cursor_peek_value(cur);
            assert(v != NULL);

            TRY(format(indent+1, w, *k, false));
            TRY(apfl_format_put_string(w, " -> "));
            TRY(format(indent+1, w, *v, true));
            TRY(apfl_format_put_string(w, "\n"));
        }
        TRY(apfl_format_put_indent(w, indent));
        TRY(apfl_format_put_string(w, "]"));
        return true;
    case VALUE_FUNC:
        TRY(apfl_format_put_string(w, "{ ... }"));
        return true;
    case VALUE_CFUNC:
        TRY(apfl_format_put_string(w, "{ native code }"));
        return true;
    case VALUE_USERDATA:
        TRY(apfl_format_put_string(w, "userdata"));
        return true;
    }

    TRY(apfl_format_put_string(w, "Unknown value ? ("));
    TRY(apfl_format_put_number(w, (int)value.type));
    TRY(apfl_format_put_string(w, ")"));
    return true;
}

enum apfl_value_type
apfl_value_type_to_abstract_type(enum value_type type)
{
    switch (type) {
    case VALUE_NIL:
        return APFL_VALUE_NIL;
    case VALUE_BOOLEAN:
        return APFL_VALUE_BOOLEAN;
    case VALUE_NUMBER:
        return APFL_VALUE_NUMBER;
    case VALUE_STRING:
    case VALUE_CONST_STRING:
        return APFL_VALUE_STRING;
    case VALUE_LIST:
        return APFL_VALUE_LIST;
    case VALUE_DICT:
        return APFL_VALUE_DICT;
    case VALUE_FUNC:
    case VALUE_CFUNC:
        return APFL_VALUE_FUNC;
    case VALUE_USERDATA:
        return APFL_VALUE_USERDATA;
    }

    assert(false);
    return 0;
}

struct function *
apfl_func_new(struct gc *gc, size_t cap, struct scope *scope)
{
    struct subfunction *subfunctions = ALLOC_LIST(gc->allocator, struct subfunction, cap);
    if (subfunctions == NULL) {
        return NULL;
    }

    struct function *function = apfl_gc_new_func(gc);
    if (function == NULL) {
        return NULL;
    }

    *function = (struct function) {
        .subfunctions = subfunctions,
        .subfunctions_len = 0,
        .subfunctions_cap = cap,
        .scope = scope,
    };

    return function;
}

bool
apfl_func_add_subfunc(struct function *function, struct instruction_list *body, struct matcher *matcher)
{
    if (function->subfunctions_len >= function->subfunctions_cap) {
        return false;
    }

    function->subfunctions[function->subfunctions_len] = (struct subfunction) {
        .body = body,
        .matcher = matcher,
    };

    function->subfunctions_len++;

    return true;
}

void
apfl_function_deinit(struct apfl_allocator allocator, struct function *function)
{
    FREE_LIST(allocator, function->subfunctions, function->subfunctions_cap);
}

struct cfunction *
apfl_cfunc_new(struct gc *gc, apfl_cfunc func, size_t nslots)
{
    struct apfl_value **slots = NULL;
    if (nslots > 0) {
        slots = ALLOC_LIST(gc->allocator, struct apfl_value *, nslots);
        if (slots == NULL) {
            return NULL;
        }
    }

    struct cfunction *cfunction = apfl_gc_new_cfunc(gc);
    if (cfunction == NULL) {
        FREE_LIST(gc->allocator, slots, nslots);
    }

    for (size_t i = 0; i < nslots; i++) {
        slots[i] = NULL;
    }

    cfunction->func = func;
    cfunction->slots = slots;
    cfunction->slots_len = nslots;

    return cfunction;
}

void
apfl_cfunction_deinit(struct apfl_allocator allocator, struct cfunction *cfunc)
{
    FREE_LIST(allocator, cfunc->slots, cfunc->slots_len);
}

struct apfl_value
apfl_value_move(struct apfl_value *src)
{
    struct apfl_value out = *src;
    src->type = VALUE_NIL;
    return out;
}

bool
apfl_value_format(struct apfl_value value, struct apfl_format_writer w)
{
    return format(0, w, value, false);
}

bool
apfl_value_print(struct apfl_value value, struct apfl_format_writer w)
{
    TRY(apfl_value_format(value, w));
    TRY(apfl_format_put_string(w, "\n"));
    return true;
}

static bool
list_eq(struct list_header *a, struct list_header *b)
{
    if (a == b) {
        return true;
    }

    if (a->len != b->len) {
        return false;
    }

    for (size_t i = 0; i < a->len; i++) {
        if (!apfl_value_eq(a->items[i], b->items[i])) {
            return false;
        }
    }
    return true;
}

static bool
dict_eq(const apfl_dict a, const apfl_dict b)
{
    if (a == b) {
        return true;
    }

    size_t total = 0;
    HASHMAP_EACH(&a->map, cur) {
        struct apfl_value *key = apfl_hashmap_cursor_peek_key(cur);
        assert(key != NULL);
        struct apfl_value *val = apfl_hashmap_cursor_peek_value(cur);
        assert(val != NULL);

        struct apfl_value *other_val = apfl_hashmap_peek(&b->map, key);
        if (other_val == NULL) {
            return false;
        }

        bool eq = apfl_value_eq(*val, *other_val);

        if (!eq) {
            return false;
        }
        total++;
    }

    return total == apfl_hashmap_count(b->map);
}

bool
apfl_value_eq(const struct apfl_value a, const struct apfl_value b)
{
    switch (a.type) {
    case VALUE_NIL:
        return b.type == VALUE_NIL && true;
    case VALUE_BOOLEAN:
        return b.type == VALUE_BOOLEAN && a.boolean == b.boolean;
    case VALUE_NUMBER:
        return b.type == VALUE_NUMBER && a.number == b.number;
    case VALUE_STRING:
    case VALUE_CONST_STRING:
        return (b.type == VALUE_STRING || b.type == VALUE_CONST_STRING)
            && apfl_string_eq(as_string_view(a), as_string_view(b));
    case VALUE_LIST:
        return b.type == VALUE_LIST && list_eq(a.list, b.list);
    case VALUE_DICT:
        return b.type == VALUE_DICT && dict_eq(a.dict, b.dict);
    case VALUE_FUNC:
        return b.type == VALUE_FUNC && a.func == b.func;
    case VALUE_CFUNC:
        return b.type == VALUE_CFUNC && a.cfunc == b.cfunc;
    case VALUE_USERDATA:
        return b.type == VALUE_USERDATA && a.userdata == b.userdata;
    }

    assert(false);
    return false;
}

#define CMP(a, b) (((a) == (b)) ? CMP_EQ : (((a) < (b)) ? CMP_LT : CMP_GT))

enum comparison_result
apfl_value_cmp(const struct apfl_value a, const struct apfl_value b)
{
    switch (a.type) {
    case VALUE_NIL:
        if (b.type != VALUE_NIL) {
            return CMP_INCOMPATIBLE_TYPES;
        }
        return CMP_EQ;
    case VALUE_BOOLEAN:
        if (b.type != VALUE_BOOLEAN) {
            return CMP_INCOMPATIBLE_TYPES;
        }
        return CMP(a.boolean ? 1 : 0, b.boolean ? 1 : 0);
    case VALUE_NUMBER:
        if (b.type != VALUE_NUMBER) {
            return CMP_INCOMPATIBLE_TYPES;
        }
        return CMP(a.number, b.number);
    case VALUE_STRING:
    case VALUE_CONST_STRING:
        if (b.type != VALUE_STRING && b.type != VALUE_CONST_STRING) {
            return CMP_INCOMPATIBLE_TYPES;
        }
        return CMP(apfl_string_cmp(as_string_view(a), as_string_view(b)), 0);
    case VALUE_LIST:
        if (b.type != VALUE_LIST) {
            return CMP_INCOMPATIBLE_TYPES;
        }
        return CMP_UNCOMPARABLE;
    case VALUE_DICT:
        if (b.type != VALUE_DICT) {
            return CMP_INCOMPATIBLE_TYPES;
        }
        return CMP_UNCOMPARABLE;
    case VALUE_FUNC:
    case VALUE_CFUNC:
        if (b.type != VALUE_FUNC && b.type != VALUE_CFUNC) {
            return CMP_INCOMPATIBLE_TYPES;
        }
        return CMP_UNCOMPARABLE;
    case VALUE_USERDATA:
        if (b.type != VALUE_USERDATA) {
            return CMP_INCOMPATIBLE_TYPES;
        }
        return CMP_UNCOMPARABLE;
    }

    assert(false);
    return CMP_INCOMPATIBLE_TYPES;
}

struct list_header *
apfl_list_new(struct gc *gc, size_t initial_cap)
{
    struct apfl_value *items = NULL;
    if (initial_cap > 0) {
        items = ALLOC_LIST(gc->allocator, struct apfl_value, initial_cap);
        if (items == NULL) {
            return NULL;
        }
    }

    struct list_header *list = apfl_gc_new_list(gc);
    if (list == NULL) {
        FREE_LIST(gc->allocator, items, initial_cap);
        return NULL;
    }

    *list = (struct list_header) {
        .items = items,
        .len = 0,
        .cap = initial_cap,
        .copy_on_write = false,
    };
    return list;
}

void
apfl_list_deinit(struct apfl_allocator allocator, struct list_header *list)
{
    FREE_LIST(allocator, list->items, list->cap);
}

bool
apfl_list_splice(
    struct gc *gc,
    struct list_header **dst_ptr,
    size_t cut_start,
    size_t cut_len,
    const struct apfl_value *other,
    size_t other_len
) {
    struct list_header *dst = *dst_ptr;

    if (!apfl_resizable_check_cut_args(dst->len, cut_start, cut_len)) {
        return false;
    }

    if (!dst->copy_on_write) {
        return apfl_resizable_splice(
            gc->allocator,
            sizeof(struct apfl_value),
            (void **)&dst->items,
            &dst->len,
            &dst->cap,
            cut_start,
            cut_len,
            other,
            other_len
        );
    }

    size_t len = dst->len - cut_len + other_len;
    struct apfl_value *items = ALLOC_LIST(gc->allocator, struct apfl_value, len);
    if (len > 0 && items == NULL) {
        return false;
    }

    struct list_header *new_list = apfl_gc_new_list(gc);
    if (new_list == NULL) {
        FREE_LIST(gc->allocator, items, len);
        return false;
    }

    // Note that we set the COW flag here, as the values now live in two places.
    for (size_t i = 0; i < cut_start; i++) {
        items[i] = apfl_value_set_cow_flag(dst->items[i]);
    }
    for (size_t i = 0; i < other_len; i++) {
        items[cut_start + i] = apfl_value_set_cow_flag(other[i]);
    }
    for (size_t i = cut_start + cut_len; i < dst->len; i++) {
        items[other_len + i - cut_len] = apfl_value_set_cow_flag(dst->items[i]);
    }

    *new_list = (struct list_header) {
        .items = items,
        .len = len,
        .cap = len,
        .copy_on_write = false,
    };

    *dst_ptr = new_list;
    return true;
}

/* Returns a dictionary for editing. Will create a copy, if neccessary.
 * *_dict must be known to the garbage collector!
 */
static struct dict_header *
dict_get_for_editing(struct gc *gc, struct dict_header **_dict)
{
    struct dict_header *dict = *_dict;
    if (!dict->copy_on_write) {
        return dict;
    }

    struct dict_header copy = {.copy_on_write = false};
    if (!apfl_hashmap_copy(&copy.map, dict->map)) {
        return NULL;
    }

    // Set the COW flags of all keys and values in the copy.
    HASHMAP_EACH(&copy.map, cur) {
        struct apfl_value *item;

        item = apfl_hashmap_cursor_peek_key(cur);
        assert(item != NULL);
        *item = apfl_value_set_cow_flag(*item);

        item = apfl_hashmap_cursor_peek_value(cur);
        assert(item != NULL);
        *item = apfl_value_set_cow_flag(*item);
    }

    dict = apfl_gc_new_dict(gc);
    if (dict == NULL) {
        apfl_hashmap_deinit(&copy.map);
        return NULL;
    }

    *dict = copy;
    *_dict = dict;

    return dict;
}

/* Set a key-value pair in a raw dictionary.
 * *_dict, k and v must all be known to the garbage collector!
 */
bool
apfl_dict_set_raw(
    struct gc *gc,
    struct dict_header **_dict,
    struct apfl_value k,
    struct apfl_value v
) {
    struct dict_header *dict = dict_get_for_editing(gc, _dict);
    if (dict == NULL) {
        return false;
    }

    return apfl_hashmap_set(&dict->map, &k, &v);
}

size_t
apfl_dict_len(struct dict_header *dict)
{
    return apfl_hashmap_count(dict->map);
}

static bool
list_get_item(struct list_header *list, size_t index, struct apfl_value *out)
{
    if (index >= list->len) {
        return false;
    }

    *out = list->items[index];
    return true;
}

static enum get_item_result
value_get_item_inner(struct apfl_value container, struct apfl_value key, struct apfl_value *out)
{
    if (container.type == VALUE_LIST) {
        if (key.type != VALUE_NUMBER) {
            return GET_ITEM_WRONG_KEY_TYPE;
        }

        return list_get_item(
            container.list,
            (size_t)key.number,
            out
        )
            ? GET_ITEM_OK
            : GET_ITEM_KEY_DOESNT_EXIST;
    } else if (container.type == VALUE_DICT) {
        return apfl_hashmap_get(
            &container.dict->map,
            &key,
            out
        )
            ? GET_ITEM_OK
            : GET_ITEM_KEY_DOESNT_EXIST;
    } else {
        return GET_ITEM_NOT_A_CONTAINER;
    }
}

enum get_item_result
apfl_value_get_item(struct apfl_value container, struct apfl_value key, struct apfl_value *out)
{
    enum get_item_result result = value_get_item_inner(container, key, out);
    if (result == GET_ITEM_OK) {
        *out = apfl_value_set_cow_flag(*out);
    }
    return result;
}

struct apfl_value
apfl_value_set_cow_flag(struct apfl_value value)
{
    switch (value.type) {
    case VALUE_LIST:
        value.list->copy_on_write = true;
        break;
    case VALUE_DICT:
        value.dict->copy_on_write = true;
        break;
    default:
        break;
    }

    return value;
}

apfl_hash
apfl_value_hash(const struct apfl_value value)
{
    apfl_hash hash = apfl_hash_fnv1a(&value.type, sizeof(enum value_type));

    struct apfl_string_view sv;

    switch (value.type) {
    case VALUE_NIL:
        return hash;
    case VALUE_BOOLEAN:
        return apfl_hash_fnv1a_add(&value.boolean, sizeof(bool), hash);
    case VALUE_NUMBER:
        return apfl_hash_fnv1a_add(&value.number, sizeof(apfl_number), hash);
    case VALUE_STRING:
    case VALUE_CONST_STRING:
        sv = as_string_view(value);
        return apfl_hash_fnv1a_add(sv.bytes, sv.len, hash);
    case VALUE_LIST:
        for (size_t i = 0; i < value.list->len; i++) {
            apfl_hash item_hash = apfl_value_hash(value.list->items[i]);
            hash = apfl_hash_fnv1a_add(&item_hash, sizeof(apfl_hash), hash);
        }
        return hash;
    case VALUE_DICT:
        // TODO: This results in all dictionaries having the same hash. Since
        //       it's rather unusual to have dictionaries as keys, this is fine
        //       for now, but should be improved nonetheless!
        return hash;
    case VALUE_FUNC:
        return apfl_hash_fnv1a_add(&value.func, sizeof(struct function *), hash);
    case VALUE_CFUNC:
        return apfl_hash_fnv1a_add(&value.cfunc, sizeof(struct cfunction *), hash);
    case VALUE_USERDATA:
        return apfl_hash_fnv1a_add(&value.userdata, sizeof(void *), hash);
    }

    assert(false);

    return hash;
}

struct gc_object *
apfl_value_get_gc_object(struct apfl_value value)
{
    switch (value.type) {
    case VALUE_NIL:
    case VALUE_BOOLEAN:
    case VALUE_NUMBER:
    case VALUE_CONST_STRING:
    case VALUE_USERDATA:
        return NULL;
    case VALUE_STRING:
        return GC_OBJECT_FROM(value.string, GC_TYPE_STRING);
    case VALUE_LIST:
        return GC_OBJECT_FROM(value.list, GC_TYPE_LIST);
    case VALUE_DICT:
        return GC_OBJECT_FROM(value.dict, GC_TYPE_DICT);
    case VALUE_FUNC:
        return GC_OBJECT_FROM(value.func, GC_TYPE_FUNC);
    case VALUE_CFUNC:
        return GC_OBJECT_FROM(value.cfunc, GC_TYPE_CFUNC);
    }

    assert(false);
    return NULL;
}

void
apfl_value_visit_gc_object(struct apfl_value value, gc_visitor cb, void *opaque)
{
    struct gc_object *child = apfl_value_get_gc_object(value);
    if (child != NULL) {
        cb(opaque, child);
    }
}

bool
apfl_value_add_as_tmproot(struct gc *gc, struct apfl_value value)
{
    struct gc_object *obj = apfl_value_get_gc_object(value);
    return obj == NULL ? true : apfl_gc_tmproot_add(gc, obj);
}

void
apfl_gc_list_traverse(struct list_header *list, gc_visitor cb, void *opaque)
{
    for (size_t i = 0; i < list->len; i++) {
        apfl_value_visit_gc_object(list->items[i], cb, opaque);
    }
}

void
apfl_gc_dict_traverse(struct dict_header *dict, gc_visitor cb, void *opaque)
{
    HASHMAP_EACH(&dict->map, cur) {
        struct apfl_value *k = apfl_hashmap_cursor_peek_key(cur);
        assert(k != NULL);
        struct apfl_value *v = apfl_hashmap_cursor_peek_value(cur);
        assert(v != NULL);

        apfl_value_visit_gc_object(*k, cb, opaque);
        apfl_value_visit_gc_object(*v, cb, opaque);
    }
}

void
apfl_gc_func_traverse(struct function* function, gc_visitor cb, void *opaque)
{
    for (size_t i = 0; i < function->subfunctions_len; i++) {
        struct subfunction *sub = &function->subfunctions[i];
        cb(opaque, GC_OBJECT_FROM(sub->body, GC_TYPE_INSTRUCTIONS));
        cb(opaque, GC_OBJECT_FROM(sub->matcher, GC_TYPE_MATCHER));
    }
    cb(opaque, GC_OBJECT_FROM(function->scope, GC_TYPE_SCOPE));
}

void
apfl_gc_cfunc_traverse(struct cfunction* cfunc, gc_visitor cb, void *opaque)
{
    for (size_t i = 0; i < cfunc->slots_len; i++) {
        if (cfunc->slots[i] != NULL) {
            cb(opaque, GC_OBJECT_FROM(cfunc->slots[i], GC_TYPE_VAR));
        }
    }
}