#include "json.h"

#include <errno.h>
#include <stdlib.h>
#include <string.h>

typedef struct {
        const char *p;
        const char *end;
        int depth;
} json_parser_t;

static json_value_t *parse_value(json_parser_t *ps);

static json_value_t *value_new(json_type_t type)
{
        json_value_t *v = calloc(1, sizeof(*v));
        if (v != NULL)
                v->type = type;
        return v;
}

void json_free(json_value_t *value)
{
        size_t i;

        if (value == NULL)
                return;
        switch (value->type) {
        case JSON_STRING:
                free(value->as.string);
                break;
        case JSON_ARRAY:
                for (i = 0; i < value->as.array.count; i++)
                        json_free(value->as.array.items[i]);
                free(value->as.array.items);
                break;
        case JSON_OBJECT:
                for (i = 0; i < value->as.object.count; i++) {
                        free(value->as.object.keys[i]);
                        json_free(value->as.object.values[i]);
                }
                free(value->as.object.keys);
                free(value->as.object.values);
                break;
        default:
                break;
        }
        free(value);
}

static void skip_ws(json_parser_t *ps)
{
        while (ps->p < ps->end) {
                char c = *ps->p;
                if (c != ' ' && c != '\t' && c != '\n' && c != '\r')
                        return;
                ps->p++;
        }
}

static bool at(const json_parser_t *ps, char c)
{
        return ps->p < ps->end && *ps->p == c;
}

static bool literal(json_parser_t *ps, const char *word)
{
        size_t n = strlen(word);

        if ((size_t)(ps->end - ps->p) < n)
                return false;
        if (memcmp(ps->p, word, n) != 0)
                return false;
        ps->p += n;
        return true;
}

static int hex_digit(char c)
{
        if (c >= '0' && c <= '9')
                return c - '0';
        if (c >= 'a' && c <= 'f')
                return c - 'a' + 10;
        if (c >= 'A' && c <= 'F')
                return c - 'A' + 10;
        return -1;
}

static bool read_hex4(json_parser_t *ps, unsigned *out)
{
        unsigned value = 0;
        int i;

        if (ps->end - ps->p < 4)
                return false;
        for (i = 0; i < 4; i++) {
                int digit = hex_digit(ps->p[i]);
                if (digit < 0)
                        return false;
                value = (value << 4) | (unsigned)digit;
        }
        ps->p += 4;
        *out = value;
        return true;
}

static size_t utf8_encode(unsigned cp, char *out)
{
        if (cp < 0x80) {
                out[0] = (char)cp;
                return 1;
        }
        if (cp < 0x800) {
                out[0] = (char)(0xC0 | (cp >> 6));
                out[1] = (char)(0x80 | (cp & 0x3F));
                return 2;
        }
        if (cp < 0x10000) {
                out[0] = (char)(0xE0 | (cp >> 12));
                out[1] = (char)(0x80 | ((cp >> 6) & 0x3F));
                out[2] = (char)(0x80 | (cp & 0x3F));
                return 3;
        }
        out[0] = (char)(0xF0 | (cp >> 18));
        out[1] = (char)(0x80 | ((cp >> 12) & 0x3F));
        out[2] = (char)(0x80 | ((cp >> 6) & 0x3F));
        out[3] = (char)(0x80 | (cp & 0x3F));
        return 4;
}

typedef struct {
        char *data;
        size_t len;
        size_t cap;
} strbuf_t;

static bool strbuf_push(strbuf_t *sb, const char *data, size_t n)
{
        if (sb->len + n + 1 > sb->cap) {
                size_t cap = sb->cap ? sb->cap : 32;
                char *grown;
                while (sb->len + n + 1 > cap)
                        cap *= 2;
                grown = realloc(sb->data, cap);
                if (grown == NULL)
                        return false;
                sb->data = grown;
                sb->cap = cap;
        }
        memcpy(sb->data + sb->len, data, n);
        sb->len += n;
        sb->data[sb->len] = '\0';
        return true;
}

static bool parse_escape(json_parser_t *ps, strbuf_t *sb)
{
        char encoded[4];
        unsigned cp;
        char c;

        if (ps->p >= ps->end)
                return false;
        c = *ps->p++;
        switch (c) {
        case '"':
                return strbuf_push(sb, "\"", 1);
        case '\\':
                return strbuf_push(sb, "\\", 1);
        case '/':
                return strbuf_push(sb, "/", 1);
        case 'b':
                return strbuf_push(sb, "\b", 1);
        case 'f':
                return strbuf_push(sb, "\f", 1);
        case 'n':
                return strbuf_push(sb, "\n", 1);
        case 'r':
                return strbuf_push(sb, "\r", 1);
        case 't':
                return strbuf_push(sb, "\t", 1);
        case 'u':
                break;
        default:
                return false;
        }

        if (!read_hex4(ps, &cp))
                return false;
        if (cp >= 0xD800 && cp <= 0xDBFF) {
                unsigned low;
                const char *save = ps->p;
                if (ps->end - ps->p >= 2 && ps->p[0] == '\\' &&
                    ps->p[1] == 'u') {
                        ps->p += 2;
                        if (!read_hex4(ps, &low))
                                return false;
                        if (low >= 0xDC00 && low <= 0xDFFF) {
                                cp = 0x10000 + ((cp - 0xD800) << 10) +
                                     (low - 0xDC00);
                        } else {
                                ps->p = save;
                                cp = 0xFFFD;
                        }
                } else {
                        cp = 0xFFFD;
                }
        } else if (cp >= 0xDC00 && cp <= 0xDFFF) {
                cp = 0xFFFD;
        }
        return strbuf_push(sb, encoded, utf8_encode(cp, encoded));
}

static char *parse_raw_string(json_parser_t *ps)
{
        strbuf_t sb = {NULL, 0, 0};

        if (!at(ps, '"'))
                return NULL;
        ps->p++;
        if (!strbuf_push(&sb, "", 0))
                return NULL;

        while (ps->p < ps->end) {
                char c = *ps->p;
                if (c == '"') {
                        ps->p++;
                        return sb.data;
                }
                if (c == '\\') {
                        ps->p++;
                        if (!parse_escape(ps, &sb))
                                break;
                        continue;
                }
                if ((unsigned char)c < 0x20)
                        break;
                ps->p++;
                if (!strbuf_push(&sb, &c, 1))
                        break;
        }
        free(sb.data);
        return NULL;
}

static json_value_t *parse_string(json_parser_t *ps)
{
        char *text = parse_raw_string(ps);
        json_value_t *v;

        if (text == NULL)
                return NULL;
        v = value_new(JSON_STRING);
        if (v == NULL) {
                free(text);
                return NULL;
        }
        v->as.string = text;
        return v;
}

static json_value_t *parse_number(json_parser_t *ps)
{
        const char *start = ps->p;
        json_value_t *v;
        char *stop = NULL;
        char token[64];
        size_t n;
        double number;

        if (at(ps, '-'))
                ps->p++;
        if (ps->p >= ps->end)
                return NULL;
        if (*ps->p == '0')
                ps->p++;
        else if (*ps->p >= '1' && *ps->p <= '9')
                while (ps->p < ps->end && *ps->p >= '0' && *ps->p <= '9')
                        ps->p++;
        else
                return NULL;
        if (at(ps, '.')) {
                ps->p++;
                if (ps->p >= ps->end || *ps->p < '0' || *ps->p > '9')
                        return NULL;
                while (ps->p < ps->end && *ps->p >= '0' && *ps->p <= '9')
                        ps->p++;
        }
        if (at(ps, 'e') || at(ps, 'E')) {
                ps->p++;
                if (at(ps, '+') || at(ps, '-'))
                        ps->p++;
                if (ps->p >= ps->end || *ps->p < '0' || *ps->p > '9')
                        return NULL;
                while (ps->p < ps->end && *ps->p >= '0' && *ps->p <= '9')
                        ps->p++;
        }

        n = (size_t)(ps->p - start);
        if (n == 0 || n >= sizeof(token))
                return NULL;
        memcpy(token, start, n);
        token[n] = '\0';

        errno = 0;
        number = strtod(token, &stop);
        if (stop != token + n)
                return NULL;

        v = value_new(JSON_NUMBER);
        if (v == NULL)
                return NULL;
        v->as.number = number;
        return v;
}

static bool grow_slots(void ***slots, size_t count, size_t *cap, size_t size)
{
        void **grown;
        size_t next;

        if (count < *cap)
                return true;
        next = *cap ? *cap * 2 : 8;
        grown = realloc(*slots, next * size);
        if (grown == NULL)
                return false;
        memset((char *)grown + count * size, 0, (next - count) * size);
        *slots = grown;
        *cap = next;
        return true;
}

static json_value_t *parse_array(json_parser_t *ps)
{
        json_value_t *v = value_new(JSON_ARRAY);
        size_t cap = 0;

        if (v == NULL)
                return NULL;
        ps->p++;
        skip_ws(ps);
        if (at(ps, ']')) {
                ps->p++;
                return v;
        }
        for (;;) {
                json_value_t *item;
                if (!grow_slots((void ***)&v->as.array.items, v->as.array.count,
                                &cap, sizeof(json_value_t *)))
                        break;
                item = parse_value(ps);
                if (item == NULL)
                        break;
                v->as.array.items[v->as.array.count++] = item;
                skip_ws(ps);
                if (at(ps, ',')) {
                        ps->p++;
                        continue;
                }
                if (at(ps, ']')) {
                        ps->p++;
                        return v;
                }
                break;
        }
        json_free(v);
        return NULL;
}

static json_value_t *parse_object(json_parser_t *ps)
{
        json_value_t *v = value_new(JSON_OBJECT);
        size_t cap_keys = 0;
        size_t cap_values = 0;

        if (v == NULL)
                return NULL;
        ps->p++;
        skip_ws(ps);
        if (at(ps, '}')) {
                ps->p++;
                return v;
        }
        for (;;) {
                char *key;
                json_value_t *item;

                skip_ws(ps);
                if (!grow_slots((void ***)&v->as.object.keys,
                                v->as.object.count, &cap_keys, sizeof(char *)))
                        break;
                if (!grow_slots((void ***)&v->as.object.values,
                                v->as.object.count, &cap_values,
                                sizeof(json_value_t *)))
                        break;
                key = parse_raw_string(ps);
                if (key == NULL)
                        break;
                skip_ws(ps);
                if (!at(ps, ':')) {
                        free(key);
                        break;
                }
                ps->p++;
                item = parse_value(ps);
                if (item == NULL) {
                        free(key);
                        break;
                }
                v->as.object.keys[v->as.object.count] = key;
                v->as.object.values[v->as.object.count] = item;
                v->as.object.count++;

                skip_ws(ps);
                if (at(ps, ',')) {
                        ps->p++;
                        continue;
                }
                if (at(ps, '}')) {
                        ps->p++;
                        return v;
                }
                break;
        }
        json_free(v);
        return NULL;
}

static json_value_t *parse_value(json_parser_t *ps)
{
        json_value_t *v;

        if (ps->depth >= JSON_MAX_DEPTH)
                return NULL;
        skip_ws(ps);
        if (ps->p >= ps->end)
                return NULL;

        switch (*ps->p) {
        case '"':
                return parse_string(ps);
        case '{':
                ps->depth++;
                v = parse_object(ps);
                ps->depth--;
                return v;
        case '[':
                ps->depth++;
                v = parse_array(ps);
                ps->depth--;
                return v;
        case 't':
                if (!literal(ps, "true"))
                        return NULL;
                v = value_new(JSON_BOOL);
                if (v != NULL)
                        v->as.boolean = true;
                return v;
        case 'f':
                if (!literal(ps, "false"))
                        return NULL;
                v = value_new(JSON_BOOL);
                if (v != NULL)
                        v->as.boolean = false;
                return v;
        case 'n':
                if (!literal(ps, "null"))
                        return NULL;
                return value_new(JSON_NULL);
        default:
                return parse_number(ps);
        }
}

json_value_t *json_parse(const char *text, size_t len)
{
        json_parser_t ps;
        json_value_t *root;

        if (text == NULL)
                return NULL;
        ps.p = text;
        ps.end = text + len;
        ps.depth = 0;

        if (len >= 3 && memcmp(text, "\xEF\xBB\xBF", 3) == 0)
                ps.p += 3;

        root = parse_value(&ps);
        if (root == NULL)
                return NULL;
        skip_ws(&ps);
        if (ps.p != ps.end) {
                json_free(root);
                return NULL;
        }
        return root;
}

const json_value_t *json_object_get(const json_value_t *object, const char *key)
{
        size_t i;

        if (object == NULL || object->type != JSON_OBJECT || key == NULL)
                return NULL;
        for (i = 0; i < object->as.object.count; i++)
                if (strcmp(object->as.object.keys[i], key) == 0)
                        return object->as.object.values[i];
        return NULL;
}

const char *json_as_string(const json_value_t *value)
{
        if (value == NULL || value->type != JSON_STRING)
                return NULL;
        return value->as.string;
}

bool json_as_double(const json_value_t *value, double *out)
{
        char *stop = NULL;
        double parsed;

        if (value == NULL)
                return false;
        if (value->type == JSON_NUMBER) {
                *out = value->as.number;
                return true;
        }
        if (value->type != JSON_STRING || value->as.string[0] == '\0')
                return false;
        errno = 0;
        parsed = strtod(value->as.string, &stop);
        if (stop == NULL || *stop != '\0' || errno == ERANGE)
                return false;
        *out = parsed;
        return true;
}

bool json_as_llong(const json_value_t *value, long long *out)
{
        double parsed;

        if (!json_as_double(value, &parsed))
                return false;
        *out = (long long)parsed;
        return true;
}