#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(x) do { if (!(x)) return false; } while (0)

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 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:
        TRY(apfl_format_put_string(w, "\""));
        TRY(apfl_format_put_string(w, *value.string));
        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;
    }

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

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

bool
apfl_value_print(struct apfl_value value, FILE *out)
{
    struct apfl_format_writer w = apfl_format_file_writer(out);
    TRY(format(0, w, value, false));
    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 true;
    }

    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)
{
    if (a.type != b.type) {
        return false;
    }

    switch (a.type) {
    case VALUE_NIL:
        return true;
    case VALUE_BOOLEAN:
        return a.boolean == b.boolean;
    case VALUE_NUMBER:
        return a.number == b.number;
    case VALUE_STRING:
        return a.string == b.string || apfl_string_eq(*a.string, *b.string);
    case VALUE_LIST:
        return list_eq(a.list, b.list);
    case VALUE_DICT:
        return dict_eq(a.dict, b.dict);
    }

    assert(false);
    return false;
}

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

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:
        sv = apfl_string_view_from(*value.string);
        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;
    }

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

    assert(false);
    return NULL;
}

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

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

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

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