#include "core_json.h"
#include "core_log.h"

typedef struct json_parser_t json_parser_t;
struct json_parser_t {
    ma_arena_t *arena;
    ma_arena_t *temp;
    char *at;
    char *end;
    i32 line;
    i32 column;
    b32 has_error;
    s8_t error;
};

typedef struct json_array_node_t json_array_node_t;
struct json_array_node_t {
    json_array_node_t *next;
    json_t *value;
};

typedef struct json_object_node_t json_object_node_t;
struct json_object_node_t {
    json_object_node_t *next;
    json_kv_t kv;
};

fn json_t *json_parse_value(json_parser_t *par);

fn json_t *json_new(json_parser_t *par, json_kind_t kind) {
    json_t *result = ma_push_type(par->arena, json_t);
    result->kind = kind;
    return result;
}

fn b32 json_at_end(json_parser_t *par) {
    b32 result = par->at >= par->end;
    return result;
}

fn void json_errorf(json_parser_t *par, const char *fmt, ...) {
    if (par->has_error) return;
    par->has_error = true;

    va_list args;
    va_start(args, fmt);
    s8_t msg = s8_vfmt(par->temp, fmt, args);
    va_end(args);

    par->error = s8_printf(par->temp, "line %d, column %d: %S", par->line, par->column, msg);
}

fn void json_advance(json_parser_t *par) {
    if (json_at_end(par)) return;
    if (par->at[0] == '\n') {
        par->line += 1;
        par->column = 1;
    } else {
        par->column += 1;
    }
    par->at += 1;
}

fn b32 json_is_ws(char c) {
    // JSON whitespace is deliberately narrower than C whitespace.
    // In particular, vertical tab/form feed are not valid JSON whitespace.
    b32 result = c == ' ' || c == '\t' || c == '\n' || c == '\r';
    return result;
}

fn void json_skip_ws(json_parser_t *par) {
    while (!json_at_end(par) && json_is_ws(par->at[0])) json_advance(par);
}

fn b32 json_match(json_parser_t *par, s8_t string) {
    if (par->end - par->at < string.len) return false;
    s8_t at = s8_make(par->at, string.len);
    if (!s8_are_equal(at, string)) return false;
    for (i64 i = 0; i < string.len; i += 1) json_advance(par);
    return true;
}

fn b32 json_expect_char(json_parser_t *par, char c) {
    if (!json_at_end(par) && par->at[0] == c) {
        json_advance(par);
        return true;
    }
    json_errorf(par, "expected '%c'", c);
    return false;
}

fn i32 json_hex_value(char c) {
    if ('0' <= c && c <= '9') return c - '0';
    if ('a' <= c && c <= 'f') return c - 'a' + 10;
    if ('A' <= c && c <= 'F') return c - 'A' + 10;
    return -1;
}

fn void json_sb_putc(sb8_t *sb, char c) {
    char *p = ma_push_type(sb->arena, char);
    *p = c;
    sb8_append(sb, s8_make(p, 1));
}

fn void json_sb_put_utf8(sb8_t *sb, u32 code) {
    // Append UTF-8 bytes for decoded \u escapes. This uses a direct one-byte
    // append instead of sb8_printf("%c") to avoid routing every byte through
    // stb_sprintf, which UBSAN reports as doing misaligned internal stores.
    if (code <= 0x7f) {
        json_sb_putc(sb, (char)code);
    } else if (code <= 0x7ff) {
        json_sb_putc(sb, (char)(0xc0 | (code >> 6)));
        json_sb_putc(sb, (char)(0x80 | (code & 0x3f)));
    } else if (code <= 0xffff) {
        json_sb_putc(sb, (char)(0xe0 | (code >> 12)));
        json_sb_putc(sb, (char)(0x80 | ((code >> 6) & 0x3f)));
        json_sb_putc(sb, (char)(0x80 | (code & 0x3f)));
    } else {
        json_sb_putc(sb, (char)(0xf0 | (code >> 18)));
        json_sb_putc(sb, (char)(0x80 | ((code >> 12) & 0x3f)));
        json_sb_putc(sb, (char)(0x80 | ((code >> 6) & 0x3f)));
        json_sb_putc(sb, (char)(0x80 | (code & 0x3f)));
    }
}

fn b32 json_parse_hex4(json_parser_t *par, u32 *code_out) {
    u32 code = 0;
    for (i32 i = 0; i < 4; i += 1) {
        if (json_at_end(par)) {
            json_errorf(par, "unterminated unicode escape");
            return false;
        }
        i32 h = json_hex_value(par->at[0]);
        if (h < 0) {
            json_errorf(par, "bad unicode escape");
            return false;
        }
        code = (code << 4) | (u32)h;
        json_advance(par);
    }
    *code_out = code;
    return true;
}

fn s8_t json_parse_string_data(json_parser_t *par) {
    if (!json_expect_char(par, '"')) return s8_null;

    sb8_t sb = {.arena = par->temp};
    while (!json_at_end(par) && par->at[0] != '"') {
        u8 c = (u8)par->at[0];
        if (c < 0x20) {
            json_errorf(par, "control character in string");
            return s8_null;
        }

        if (c == '\\') {
            json_advance(par);
            if (json_at_end(par)) {
                json_errorf(par, "unterminated escape sequence");
                return s8_null;
            }

            switch (par->at[0]) {
            case '"': json_sb_putc(&sb, '"'); json_advance(par); break;
            case '\\': json_sb_putc(&sb, '\\'); json_advance(par); break;
            case '/':  json_sb_putc(&sb, '/');  json_advance(par); break;
            case 'b':  json_sb_putc(&sb, '\b'); json_advance(par); break;
            case 'f':  json_sb_putc(&sb, '\f'); json_advance(par); break;
            case 'n':  json_sb_putc(&sb, '\n'); json_advance(par); break;
            case 'r':  json_sb_putc(&sb, '\r'); json_advance(par); break;
            case 't':  json_sb_putc(&sb, '\t'); json_advance(par); break;
            case 'u': {
                json_advance(par);
                u32 code = 0;
                if (!json_parse_hex4(par, &code)) return s8_null;
                if (0xd800 <= code && code <= 0xdbff) {
                    if (par->end - par->at < 6 || par->at[0] != '\\' || par->at[1] != 'u') {
                        json_errorf(par, "expected low surrogate after high surrogate");
                        return s8_null;
                    }
                    json_advance(par);
                    json_advance(par);
                    u32 low = 0;
                    if (!json_parse_hex4(par, &low)) return s8_null;
                    if (!(0xdc00 <= low && low <= 0xdfff)) {
                        json_errorf(par, "bad low surrogate");
                        return s8_null;
                    }
                    code = 0x10000 + (((code - 0xd800) << 10) | (low - 0xdc00));
                } else if (0xdc00 <= code && code <= 0xdfff) {
                    json_errorf(par, "unexpected low surrogate");
                    return s8_null;
                }
                json_sb_put_utf8(&sb, code);
            } break;
            default: {
                json_errorf(par, "bad escape sequence");
                return s8_null;
            } break;
            }
        } else {
            json_sb_putc(&sb, (char)c);
            json_advance(par);
        }
    }

    if (!json_expect_char(par, '"')) return s8_null;
    s8_t temp = sb8_merge(par->temp, &sb);
    s8_t result = s8_copy(par->arena, temp);
    return result;
}

fn b32 json_is_number_delim(char c) {
    b32 result = c == 0 || char_is_whitespace(c) || c == ',' || c == ']' || c == '}';
    return result;
}

fn json_t *json_parse_number(json_parser_t *par) {
    char *first = par->at;

    if (!json_at_end(par) && par->at[0] == '-') json_advance(par);

    if (json_at_end(par)) {
        json_errorf(par, "expected number");
        return NULL;
    }

    if (par->at[0] == '0') {
        json_advance(par);
        if (!json_at_end(par) && char_is_digit(par->at[0])) {
            json_errorf(par, "leading zero in number");
            return NULL;
        }
    } else if (char_is_digit(par->at[0])) {
        while (!json_at_end(par) && char_is_digit(par->at[0])) json_advance(par);
    } else {
        json_errorf(par, "expected number");
        return NULL;
    }

    if (!json_at_end(par) && par->at[0] == '.') {
        json_advance(par);
        if (json_at_end(par) || !char_is_digit(par->at[0])) {
            json_errorf(par, "expected digit after decimal point");
            return NULL;
        }
        while (!json_at_end(par) && char_is_digit(par->at[0])) json_advance(par);
    }

    if (!json_at_end(par) && (par->at[0] == 'e' || par->at[0] == 'E')) {
        json_advance(par);
        if (!json_at_end(par) && (par->at[0] == '+' || par->at[0] == '-')) json_advance(par);
        if (json_at_end(par) || !char_is_digit(par->at[0])) {
            json_errorf(par, "expected digit in exponent");
            return NULL;
        }
        while (!json_at_end(par) && char_is_digit(par->at[0])) json_advance(par);
    }

    if (!json_is_number_delim(json_at_end(par) ? 0 : par->at[0])) {
        json_errorf(par, "invalid character after number");
        return NULL;
    }

    json_t *result = json_new(par, json_kind_number);
    result->number = f64_from_s8(s8_from_range(first, par->at));
    return result;
}

fn i32 json_parse_next(json_parser_t *par, char close, s8_t what) {
    json_skip_ws(par);
    if (!json_at_end(par) && par->at[0] == close) {
        json_advance(par);
        return 0;
    }
    if (json_at_end(par) || par->at[0] != ',') {
        json_errorf(par, "expected ',' or '%c' in %S", close, what);
        return -1;
    }
    json_advance(par);
    json_skip_ws(par);
    return 1;
}

fn json_t *json_parse_array(json_parser_t *par) {
    json_t *result = json_new(par, json_kind_array);
    // Collect elements as scratch linked-list nodes while the final count is
    // unknown, then copy once to a tight user-arena array.
    json_advance(par);
    json_skip_ws(par);

    json_array_node_t *first = NULL;
    json_array_node_t *last = NULL;
    i32 count = 0;

    if (!json_at_end(par) && par->at[0] == ']') {
        json_advance(par);
    } else {
        for (;;) {
            json_t *value = json_parse_value(par);
            if (par->has_error) return result;

            json_array_node_t *node = ma_push_type(par->temp, json_array_node_t);
            node->value = value;
            SLLQ_APPEND(first, last, node);
            count += 1;

            i32 next = json_parse_next(par, ']', s8("array"));
            if (next <= 0) break;
        }
    }

    result->array.len = count;
    result->array.data = ma_push_array(par->arena, json_t *, count);
    i32 idx = 0;
    for (json_array_node_t *it = first; it; it = it->next) {
        result->array.data[idx++] = it->value;
    }
    return result;
}

fn json_t *json_parse_object(json_parser_t *par) {
    json_t *result = json_new(par, json_kind_object);
    // Collect members in scratch first. After the object is complete we can
    // validate duplicate keys and copy to one contiguous json_kv_t array.
    json_advance(par);
    json_skip_ws(par);

    json_object_node_t *first = NULL;
    json_object_node_t *last = NULL;
    i32 count = 0;

    if (!json_at_end(par) && par->at[0] == '}') {
        json_advance(par);
    } else {
        for (;;) {
            if (json_at_end(par) || par->at[0] != '"') {
                json_errorf(par, "expected object key string");
                return result;
            }
            s8_t key = json_parse_string_data(par);
            if (par->has_error) return result;

            json_skip_ws(par);
            if (!json_expect_char(par, ':')) return result;
            json_skip_ws(par);

            json_t *value = json_parse_value(par);
            if (par->has_error) return result;

            json_object_node_t *node = ma_push_type(par->temp, json_object_node_t);
            node->kv = (json_kv_t){.key = key, .value = value};
            SLLQ_APPEND(first, last, node);
            count += 1;

            i32 next = json_parse_next(par, '}', s8("object"));
            if (next <= 0) break;
        }
    }

    for (json_object_node_t *a = first; a; a = a->next) {
        for (json_object_node_t *b = a->next; b; b = b->next) {
            if (s8_are_equal(a->kv.key, b->kv.key)) {
                json_errorf(par, "duplicate object key: %S", a->kv.key);
                return result;
            }
        }
    }

    result->object.len = count;
    result->object.data = ma_push_array(par->arena, json_kv_t, count);
    i32 idx = 0;
    for (json_object_node_t *it = first; it; it = it->next) {
        result->object.data[idx++] = it->kv;
    }
    return result;
}

fn json_t *json_parse_value(json_parser_t *par) {
    json_skip_ws(par);
    if (json_at_end(par)) {
        json_errorf(par, "expected JSON value");
        return NULL;
    }

    switch (par->at[0]) {
    case 'n': {
        if (json_match(par, s8("null"))) return json_new(par, json_kind_null);
    } break;
    case 't': {
        if (json_match(par, s8("true"))) {
            json_t *result = json_new(par, json_kind_bool);
            result->boolean = true;
            return result;
        }
    } break;
    case 'f': {
        if (json_match(par, s8("false"))) {
            json_t *result = json_new(par, json_kind_bool);
            result->boolean = false;
            return result;
        }
    } break;
    case '"': {
        json_t *result = json_new(par, json_kind_string);
        result->string = json_parse_string_data(par);
        return par->has_error ? NULL : result;
    } break;
    case '[': return json_parse_array(par);
    case '{': return json_parse_object(par);
    default: {
        if (par->at[0] == '-' || char_is_digit(par->at[0])) return json_parse_number(par);
    } break;
    }

    json_errorf(par, "expected JSON value");
    return NULL;
}

fn json_t *json_parse(ma_arena_t *arena, s8_t stream) {
    ma_temp_t temp = ma_begin_scratch1(arena);
    json_parser_t par = {.arena = arena, .temp = temp.arena, .at = stream.str, .end = stream.str + stream.len, .line = 1, .column = 1};

    // Accept an optional UTF-8 BOM at the beginning of a file/buffer.
    if (stream.len >= 3 && (u8)stream.str[0] == 0xef && (u8)stream.str[1] == 0xbb && (u8)stream.str[2] == 0xbf) {
        json_advance(&par);
        json_advance(&par);
        json_advance(&par);
    }

    json_t *result = json_parse_value(&par);
    if (!par.has_error) {
        json_skip_ws(&par);
        if (!json_at_end(&par)) json_errorf(&par, "unexpected trailing data");
    }

    if (par.has_error) {
        errorf("json parse error: %S", par.error);
        result = NULL;
    }

    ma_end_scratch(temp);
    return result;
}

fn void json_serialize_string(sb8_t *sb, s8_t string) {
    sb8_append(sb, s8("\""));
    for (i64 i = 0; i < string.len; i += 1) {
        u8 c = (u8)string.str[i];
        switch (c) {
        case '"': sb8_append(sb, s8("\\\"")); break;
        case '\\': sb8_append(sb, s8("\\\\")); break;
        case '\b': sb8_append(sb, s8("\\b")); break;
        case '\f': sb8_append(sb, s8("\\f")); break;
        case '\n': sb8_append(sb, s8("\\n")); break;
        case '\r': sb8_append(sb, s8("\\r")); break;
        case '\t': sb8_append(sb, s8("\\t")); break;
        default: {
            if (c < 0x20) sb8_printf(sb, "\\u%04x", c);
            else json_sb_putc(sb, (char)c);
        } break;
        }
    }
    sb8_append(sb, s8("\""));
}

fn void json_indent(sb8_t *sb, s8_t new_line, s8_t indent, i32 depth) {
    if (!new_line.str) return;
    sb8_append(sb, new_line);
    for (i32 i = 0; i < depth; i += 1) sb8_append(sb, indent);
}

fn void json_serialize_into(sb8_t *sb, json_t *n, s8_t new_line, s8_t indent, i32 depth) {
    if (!n) {
        sb8_append(sb, s8("null"));
        return;
    }

    switch (n->kind) {
    case json_kind_null: sb8_append(sb, s8("null")); break;
    case json_kind_bool: sb8_append(sb, n->boolean ? s8("true") : s8("false")); break;
    case json_kind_number: sb8_printf(sb, "%.17g", n->number); break;
    case json_kind_string: json_serialize_string(sb, n->string); break;
    case json_kind_array: {
        sb8_append(sb, s8("["));
        for (i32 i = 0; i < n->array.len; i += 1) {
            if (i) sb8_append(sb, s8(","));
            json_indent(sb, new_line, indent, depth + 1);
            json_serialize_into(sb, n->array.data[i], new_line, indent, depth + 1);
        }
        if (n->array.len) json_indent(sb, new_line, indent, depth);
        sb8_append(sb, s8("]"));
    } break;
    case json_kind_object: {
        sb8_append(sb, s8("{"));
        for (i32 i = 0; i < n->object.len; i += 1) {
            if (i) sb8_append(sb, s8(","));
            json_indent(sb, new_line, indent, depth + 1);
            json_serialize_string(sb, n->object.data[i].key);
            sb8_append(sb, new_line.str ? s8(": ") : s8(":"));
            json_serialize_into(sb, n->object.data[i].value, new_line, indent, depth + 1);
        }
        if (n->object.len) json_indent(sb, new_line, indent, depth);
        sb8_append(sb, s8("}"));
    } break;
    }
}

fn s8_t json_serialize_compact(ma_arena_t *arena, json_t *n) {
    ma_temp_t temp = ma_begin_scratch1(arena);
    sb8_t sb = {.arena = temp.arena};
    json_serialize_into(&sb, n, s8_null, s8_null, 0);
    s8_t result = sb8_merge(arena, &sb);
    ma_end_scratch(temp);
    return result;
}

fn s8_t json_serialize_pretty(ma_arena_t *arena, json_t *n, s8_t new_line, s8_t indent) {
    ma_temp_t temp = ma_begin_scratch1(arena);
    sb8_t sb = {.arena = temp.arena};
    json_serialize_into(&sb, n, new_line.str ? new_line : s8("\n"), indent.str ? indent : s8("  "), 0);
    s8_t result = sb8_merge(arena, &sb);
    ma_end_scratch(temp);
    return result;
}

fn json_t *json_object_get(json_t *n, s8_t key) {
    if (!n || n->kind != json_kind_object) return NULL;
    for (i32 i = 0; i < n->object.len; i += 1) {
        if (s8_are_equal(n->object.data[i].key, key)) return n->object.data[i].value;
    }
    return NULL;
}

fn json_t *json_array_get(json_t *n, i32 idx) {
    if (!n || n->kind != json_kind_array || idx < 0 || idx >= n->array.len) return NULL;
    return n->array.data[idx];
}

fn json_t *jqv(json_t *n, ma_arena_t *temp, const char *fmt, va_list args) {
    // Query parsing only needs transient formatting storage; returned json_t is
    // from the original tree, so callers should not have to pass an arena.
    s8_t path = s8_vfmt(temp, fmt, args);
    s8_t rest = path;
    json_t *cur = n;

    while (cur && rest.len > 0) {
        if (rest.str[0] == '.') {
            rest = s8_skip(rest, 1);
            continue;
        }

        if (rest.str[0] == '[') {
            rest = s8_skip(rest, 1);
            i64 idx = 0;
            b32 any = false;
            while (rest.len > 0 && char_is_digit(rest.str[0])) {
                any = true;
                idx = idx * 10 + (rest.str[0] - '0');
                rest = s8_skip(rest, 1);
            }
            if (!any || rest.len == 0 || rest.str[0] != ']') return NULL;
            rest = s8_skip(rest, 1);
            cur = json_array_get(cur, (i32)idx);
        } else {
            char *key_begin = rest.str;
            while (rest.len > 0 && rest.str[0] != '.' && rest.str[0] != '[') rest = s8_skip(rest, 1);
            cur = json_object_get(cur, s8_from_range(key_begin, rest.str));
        }
    }

    return cur;
}

fn json_t *jq(json_t *n, const char *fmt, ...) {
    ma_temp_t temp = ma_begin_scratch();
    va_list args;
    va_start(args, fmt);
    json_t *result = jqv(n, temp.arena, fmt, args);
    va_end(args);
    ma_end_scratch(temp);
    return result;
}

fn f64 jqn(json_t *n, const char *fmt, ...) {
    ma_temp_t temp = ma_begin_scratch();
    va_list args;
    va_start(args, fmt);
    json_t *result = jqv(n, temp.arena, fmt, args);
    va_end(args);
    ma_end_scratch(temp);
    return result && result->kind == json_kind_number ? result->number : 0;
}

fn s8_t jqs(json_t *n, const char *fmt, ...) {
    ma_temp_t temp = ma_begin_scratch();
    va_list args;
    va_start(args, fmt);
    json_t *result = jqv(n, temp.arena, fmt, args);
    va_end(args);
    ma_end_scratch(temp);
    return result && result->kind == json_kind_string ? result->string : s8_null;
}

fn i32 jql(json_t *n, const char *fmt, ...) {
    ma_temp_t temp = ma_begin_scratch();
    va_list args;
    va_start(args, fmt);
    json_t *result = jqv(n, temp.arena, fmt, args);
    va_end(args);
    ma_end_scratch(temp);
    if (!result) return 0;
    if (result->kind == json_kind_array) return result->array.len;
    if (result->kind == json_kind_object) return result->object.len;
    return 0;
}

fn b32 json_equal(json_t *a, json_t *b) {
    if (a == b) return true;
    if (!a || !b || a->kind != b->kind) return false;

    switch (a->kind) {
    case json_kind_null: return true;
    case json_kind_bool: return a->boolean == b->boolean;
    case json_kind_number: return a->number == b->number;
    case json_kind_string: return s8_are_equal(a->string, b->string);
    case json_kind_array: {
        if (a->array.len != b->array.len) return false;
        for (i32 i = 0; i < a->array.len; i += 1) {
            if (!json_equal(a->array.data[i], b->array.data[i])) return false;
        }
        return true;
    } break;
    case json_kind_object: {
        if (a->object.len != b->object.len) return false;
        for (i32 i = 0; i < a->object.len; i += 1) {
            if (!s8_are_equal(a->object.data[i].key, b->object.data[i].key)) return false;
            if (!json_equal(a->object.data[i].value, b->object.data[i].value)) return false;
        }
        return true;
    } break;
    }

    return false;
}