#include "json.h"

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

#define JSMN_ERROR_NOMEM -1
#define JSMN_ERROR_INVAL -2
#define JSMN_ERROR_PART -3

static jsmntok_t *jsmn_alloc_token(jsmn_parser *parser, jsmntok_t *tokens, size_t num_tokens) {
    jsmntok_t *tok;

    if (parser->toknext >= num_tokens) {
        return NULL;
    }

    tok = &tokens[parser->toknext++];
    tok->start = -1;
    tok->end = -1;
    tok->size = 0;
    tok->parent = -1;
    tok->type = JSMN_UNDEFINED;
    return tok;
}

static void jsmn_fill_token(jsmntok_t *token, jsmntype_t type, int start, int end) {
    token->type = type;
    token->start = start;
    token->end = end;
    token->size = 0;
}

void jsmn_init(jsmn_parser *parser) {
    parser->pos = 0;
    parser->toknext = 0;
    parser->toksuper = -1;
}

static int jsmn_parse_primitive(jsmn_parser *parser, const char *js, size_t len, jsmntok_t *tokens, size_t num_tokens) {
    int start = (int)parser->pos;

    for (; parser->pos < len; parser->pos++) {
        char c = js[parser->pos];
        if (c == ':' || c == '\t' || c == '\r' || c == '\n' || c == ' ' || c == ',' || c == ']' || c == '}') {
            jsmntok_t *token = jsmn_alloc_token(parser, tokens, num_tokens);
            if (token == NULL) {
                parser->pos = (unsigned int)start;
                return JSMN_ERROR_NOMEM;
            }
            jsmn_fill_token(token, JSMN_PRIMITIVE, start, (int)parser->pos);
            token->parent = parser->toksuper;
            parser->pos--;
            return 0;
        }

        if (c < 32 || c >= 127) {
            parser->pos = (unsigned int)start;
            return JSMN_ERROR_INVAL;
        }
    }

    parser->pos = (unsigned int)start;
    return JSMN_ERROR_PART;
}

static int jsmn_parse_string(jsmn_parser *parser, const char *js, size_t len, jsmntok_t *tokens, size_t num_tokens) {
    int start = (int)parser->pos;

    parser->pos++;

    for (; parser->pos < len; parser->pos++) {
        char c = js[parser->pos];

        if (c == '"') {
            jsmntok_t *token = jsmn_alloc_token(parser, tokens, num_tokens);
            if (token == NULL) {
                parser->pos = (unsigned int)start;
                return JSMN_ERROR_NOMEM;
            }
            jsmn_fill_token(token, JSMN_STRING, start + 1, (int)parser->pos);
            token->parent = parser->toksuper;
            return 0;
        }

        if (c == '\\' && parser->pos + 1 < len) {
            char esc = js[parser->pos + 1];

            if (esc == '"' || esc == '/' || esc == '\\' || esc == 'b' || esc == 'f' || esc == 'r' || esc == 'n' || esc == 't') {
                parser->pos++;
                continue;
            }

            if (esc == 'u') {
                int i;
                parser->pos++;
                for (i = 0; i < 4 && parser->pos + 1 < len; i++) {
                    char hc = js[parser->pos + 1];
                    if (!((hc >= '0' && hc <= '9') || (hc >= 'A' && hc <= 'F') || (hc >= 'a' && hc <= 'f'))) {
                        parser->pos = (unsigned int)start;
                        return JSMN_ERROR_INVAL;
                    }
                    parser->pos++;
                }
                continue;
            }

            parser->pos = (unsigned int)start;
            return JSMN_ERROR_INVAL;
        }
    }

    parser->pos = (unsigned int)start;
    return JSMN_ERROR_PART;
}

int jsmn_parse(jsmn_parser *parser, const char *js, size_t len, jsmntok_t *tokens, unsigned int num_tokens) {
    int r;
    unsigned int i;

    for (; parser->pos < len; parser->pos++) {
        char c = js[parser->pos];

        switch (c) {
            case '{':
            case '[': {
                jsmntype_t type = (c == '{') ? JSMN_OBJECT : JSMN_ARRAY;
                jsmntok_t *token = jsmn_alloc_token(parser, tokens, num_tokens);

                if (token == NULL) {
                    return JSMN_ERROR_NOMEM;
                }

                if (parser->toksuper != -1) {
                    tokens[parser->toksuper].size++;
                    token->parent = parser->toksuper;
                }

                jsmn_fill_token(token, type, (int)parser->pos, -1);
                parser->toksuper = (int)(parser->toknext - 1);
                break;
            }

            case '}':
            case ']': {
                jsmntype_t type = (c == '}') ? JSMN_OBJECT : JSMN_ARRAY;

                for (i = parser->toknext; i > 0; i--) {
                    jsmntok_t *token = &tokens[i - 1];
                    if (token->start != -1 && token->end == -1) {
                        if (token->type != type) {
                            return JSMN_ERROR_INVAL;
                        }
                        token->end = (int)parser->pos + 1;
                        parser->toksuper = token->parent;
                        break;
                    }
                }

                if (i == 0) {
                    return JSMN_ERROR_INVAL;
                }
                break;
            }

            case '"':
                r = jsmn_parse_string(parser, js, len, tokens, num_tokens);
                if (r < 0) {
                    return r;
                }
                if (parser->toksuper != -1) {
                    tokens[parser->toksuper].size++;
                }
                break;

            case '\t':
            case '\r':
            case '\n':
            case ' ':
            case ':':
            case ',':
                break;

            default:
                r = jsmn_parse_primitive(parser, js, len, tokens, num_tokens);
                if (r < 0) {
                    return r;
                }
                if (parser->toksuper != -1) {
                    tokens[parser->toksuper].size++;
                }
                break;
        }
    }

    for (i = parser->toknext; i > 0; i--) {
        if (tokens[i - 1].start != -1 && tokens[i - 1].end == -1) {
            return JSMN_ERROR_PART;
        }
    }

    return (int)parser->toknext;
}

int json_doc_parse(JsonDoc *doc, const char *json) {
    jsmn_parser p;
    int r;
    int cap = 1024;

    doc->json = json;
    doc->tokens = NULL;
    doc->token_count = 0;

    while (1) {
        doc->tokens = (jsmntok_t *)malloc((size_t)cap * sizeof(jsmntok_t));
        if (doc->tokens == NULL) {
            return 0;
        }

        jsmn_init(&p);
        r = jsmn_parse(&p, json, strlen(json), doc->tokens, (unsigned int)cap);
        if (r >= 0) {
            doc->token_count = r;
            return 1;
        }

        free(doc->tokens);
        doc->tokens = NULL;

        if (r == JSMN_ERROR_NOMEM) {
            cap *= 2;
            if (cap > (1 << 21)) {
                return 0;
            }
            continue;
        }

        return 0;
    }
}

void json_doc_free(JsonDoc *doc) {
    free(doc->tokens);
    doc->tokens = NULL;
    doc->token_count = 0;
    doc->json = NULL;
}

int json_token_streq(const JsonDoc *doc, int tok_idx, const char *s) {
    const jsmntok_t *tok;
    size_t len;

    if (tok_idx < 0 || tok_idx >= doc->token_count) {
        return 0;
    }

    tok = &doc->tokens[tok_idx];
    if (tok->type != JSMN_STRING && tok->type != JSMN_PRIMITIVE) {
        return 0;
    }

    len = (size_t)(tok->end - tok->start);
    if (len != strlen(s)) {
        return 0;
    }

    return strncmp(doc->json + tok->start, s, len) == 0;
}

int json_skip_token(const JsonDoc *doc, int tok_idx) {
    int i;
    int next;

    if (tok_idx < 0 || tok_idx >= doc->token_count) {
        return -1;
    }

    if (doc->tokens[tok_idx].type == JSMN_STRING || doc->tokens[tok_idx].type == JSMN_PRIMITIVE) {
        return tok_idx + 1;
    }

    if (doc->tokens[tok_idx].type == JSMN_ARRAY) {
        next = tok_idx + 1;
        for (i = 0; i < doc->tokens[tok_idx].size; i++) {
            next = json_skip_token(doc, next);
            if (next < 0) {
                return -1;
            }
        }
        return next;
    }

    if (doc->tokens[tok_idx].type == JSMN_OBJECT) {
        int pairs = doc->tokens[tok_idx].size / 2;
        next = tok_idx + 1;
        for (i = 0; i < pairs; i++) {
            next = json_skip_token(doc, next);
            if (next < 0) {
                return -1;
            }
            next = json_skip_token(doc, next);
            if (next < 0) {
                return -1;
            }
        }
        return next;
    }

    return -1;
}

int json_object_get(const JsonDoc *doc, int obj_idx, const char *key) {
    int i;
    int cur;
    int pairs;

    if (obj_idx < 0 || obj_idx >= doc->token_count) {
        return -1;
    }

    if (doc->tokens[obj_idx].type != JSMN_OBJECT) {
        return -1;
    }

    pairs = doc->tokens[obj_idx].size / 2;
    cur = obj_idx + 1;
    for (i = 0; i < pairs; i++) {
        int key_idx = cur;
        int value_idx;

        cur = json_skip_token(doc, cur);
        if (cur < 0) {
            return -1;
        }

        value_idx = cur;
        cur = json_skip_token(doc, cur);
        if (cur < 0) {
            return -1;
        }

        if (json_token_streq(doc, key_idx, key)) {
            return value_idx;
        }
    }

    return -1;
}

int json_array_get(const JsonDoc *doc, int arr_idx, int elem_index) {
    int i;
    int cur;

    if (arr_idx < 0 || arr_idx >= doc->token_count) {
        return -1;
    }

    if (doc->tokens[arr_idx].type != JSMN_ARRAY) {
        return -1;
    }

    if (elem_index < 0 || elem_index >= doc->tokens[arr_idx].size) {
        return -1;
    }

    cur = arr_idx + 1;
    for (i = 0; i < elem_index; i++) {
        cur = json_skip_token(doc, cur);
        if (cur < 0) {
            return -1;
        }
    }

    return cur;
}

char *json_token_strdup(const JsonDoc *doc, int tok_idx) {
    int start;
    int end;
    size_t len;
    char *out;

    if (tok_idx < 0 || tok_idx >= doc->token_count) {
        return NULL;
    }

    start = doc->tokens[tok_idx].start;
    end = doc->tokens[tok_idx].end;
    if (start < 0 || end < start) {
        return NULL;
    }

    len = (size_t)(end - start);
    out = (char *)malloc(len + 1);
    if (out == NULL) {
        return NULL;
    }

    memcpy(out, doc->json + start, len);
    out[len] = '\0';
    return out;
}