#include "json_collections.h"
#include "json.h"
#include <assert.h>
#include <stdbool.h>
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

struct json_value {
    enum json_type ty;
    union {
        int64_t int_val;
        double float_val;
        char *string_val;
        struct smallarray array_vals;
        struct hashmap object_fields;
    };
};

#define shallow_int_min (-144115188075855872ll)
#define shallow_int_max 144115188075855871ll

static bool is_shallow(struct json_value const *ptr)
{
    return ((size_t)ptr >> 1 & 1) != 0;
}

static enum json_type shallow_ty(struct json_value const *ptr)
{
    return (enum json_type)((size_t)ptr >> 2 & 0xf);
}

static void set_shallow_ty(struct json_value **ptr, enum json_type ty)
{
    *(size_t *)ptr |= (ty & 0xf) << 2;
}

static size_t shallow_val(struct json_value const *ptr)
{
    return (size_t)ptr >> 6;
}

static void set_shallow_val(struct json_value **ptr, size_t val)
{
    *(size_t *)ptr |= val << 6;
}

static struct json_value *new_shallow(enum json_type ty)
{
    struct json_value *shallow_val = (struct json_value *)0x2;
    set_shallow_ty(&shallow_val, ty);
    return shallow_val;
}

static struct json_value *new_deep(enum json_type ty)
{
    struct json_value *v = malloc(sizeof(struct json_value));
    *v = (struct json_value) { .ty = ty };
    switch (ty) {
        case json_null:
        case json_false:
        case json_true:
        case json_int:
        case json_float:
        case json_string:
            break;
        case json_array:
            smallarray_construct(&v->array_vals);
            break;
        case json_object:
            hashmap_construct(&v->object_fields);
            break;
    }
    return v;
}

struct json_value *json_new(enum json_type ty)
{
    switch (ty) {
        case json_null:
        case json_false:
        case json_true:
        case json_int:
        case json_float:
            return new_shallow(ty);
        case json_string:
        case json_array:
        case json_object:
            return new_deep(ty);
    }
    assert(false);
}

void json_free(struct json_value *v)
{
    if (is_shallow(v))
        return;

    switch (v->ty) {
        case json_null:
        case json_false:
        case json_true:
        case json_int:
        case json_float:
            break;
        case json_string:
            if (v->string_val)
                free(v->string_val);
            break;
        case json_array:
            smallarray_destroy(&v->array_vals);
            break;
        case json_object:
            hashmap_destroy(&v->object_fields);
            break;
    }
    free(v);
}

bool json_is(struct json_value const *value, enum json_type type)
{
    if (is_shallow(value))
        return shallow_ty(value) == type;

    return value->ty == type;
}

enum json_type json_get_type(struct json_value const *value)
{
    if (is_shallow(value))
        return shallow_ty(value);

    return value->ty;
}

bool json_get_bool(struct json_value const *value)
{
    enum json_type ty;

    if (is_shallow(value)) {
        ty = shallow_ty(value);
    } else {
        ty = value->ty;
    }

    assert(ty == json_true || ty == json_false);
    return ty == json_true;
}
int64_t json_get_int(struct json_value const *value)
{
    if (is_shallow(value)) {
        assert(shallow_ty(value) == json_int);
        return (int64_t)shallow_val(value);
    }

    assert(value->ty == json_int);
    return value->int_val;
}
double json_get_float(struct json_value const *value)
{
    if (is_shallow(value)) {
        assert(shallow_ty(value) == json_float);
        size_t raw_val = shallow_val(value) << 6;
        return *(double *)&raw_val;
    }

    assert(value->ty == json_float);
    return value->float_val;
}
char const *json_get_string(struct json_value const *value)
{
    assert(value->ty == json_string);
    return value->string_val;
}

void json_set_null(struct json_value **value)
{
    if (is_shallow(*value)) {
        set_shallow_ty(value, json_null);
        return;
    }

    (*value)->ty = json_null;
}

void json_set_bool(struct json_value **value, bool val)
{
    if (is_shallow(*value)) {
        set_shallow_ty(value, val ? json_true : json_false);
    }

    (*value)->ty = val ? json_true : json_false;
}
void json_set_int(struct json_value **value, int64_t val)
{
    if (is_shallow(*value)) {
        if (val >= shallow_int_min && val <= shallow_int_max) {
            set_shallow_ty(value, json_int);
            set_shallow_val(value, (size_t)val);
            return;
        } else {
            *value = new_deep(json_int);
        }
    }

    (*value)->ty = json_int;
    (*value)->int_val = val;
}
void json_set_float(struct json_value **value, double val)
{
    if (is_shallow(*value)) {
        size_t raw_val = *(size_t *)&val;
        if ((raw_val & 0x7f) == 0) {
            set_shallow_ty(value, json_float);
            set_shallow_val(value, raw_val >> 6);
            return;
        } else {
            *value = new_deep(json_float);
        }
    }

    (*value)->ty = json_float;
    (*value)->float_val = val;
}
void json_set_string(struct json_value *value, char *val)
{
    value->ty = json_string;
    value->string_val = val;
}

struct json_value *json_new_int(int64_t val)
{
    struct json_value *node = json_new(json_int);
    json_set_int(&node, val);
    return node;
}

struct json_value *json_new_float(double val)
{
    struct json_value *node = json_new(json_float);
    json_set_float(&node, val);
    return node;
}

size_t json_array_count(struct json_value const *array)
{
    assert(array->ty == json_array);
    return smallarray_count(&array->array_vals);
}

struct json_value *json_idx(struct json_value *array, size_t idx)
{
    assert(array->ty == json_array);
    if (idx >= smallarray_count(&array->array_vals))
        return NULL;
    return smallarray_get(&array->array_vals, idx);
}

struct json_value const *json_idx_const(
    struct json_value const *array, size_t idx)
{
    assert(array->ty == json_array);
    if (idx >= smallarray_count(&array->array_vals))
        return NULL;
    return smallarray_get_const(&array->array_vals, idx);
}

void json_push(struct json_value *array, struct json_value *value)
{
    assert(array->ty == json_array);
    smallarray_push(&array->array_vals, value);
}

struct json_value *json_key(struct json_value *object, char const *key)
{
    assert(object->ty == json_object);
    return hashmap_get(&object->object_fields, key);
}

struct json_value *json_key_sized(
    struct json_value *object, char const *key, size_t key_size)
{
    assert(object->ty == json_object);
    return hashmap_get_sized(&object->object_fields, key, key_size);
}

struct json_value const *json_key_hash_const(
    struct json_value const *object, size_t hash)
{
    assert(object->ty == json_object);
    return hashmap_get_hash_const(&object->object_fields, hash);
}

void json_set(
    struct json_value *object, char const *key, struct json_value *value)
{
    assert(object->ty == json_object);
    hashmap_set(&object->object_fields, key, value);
}

void json_set_sized(struct json_value *object,
    char const *key,
    size_t key_size,
    struct json_value *value)
{
    assert(object->ty == json_object);
    hashmap_set_sized(&object->object_fields, key, key_size, value);
}

void json_object_keys(struct json_value const *object,
    struct hash_key_entry const **keys,
    size_t *count)
{
    assert(object->ty == json_object);
    *keys = object->object_fields.keys;
    *count = object->object_fields.keys_count;
}