#include <lcq/json.h>

enum json_value_buf_type {
    JSON_VALUE_BUF_TYPE_TERMINATOR=0x00,
    JSON_VALUE_BUF_TYPE_NUMBER_ONEBYTE_POS=0x10,
    JSON_VALUE_BUF_TYPE_NUMBER_MULTIBYTE_POS=0x20,
    JSON_VALUE_BUF_TYPE_NUMBER_ONEBYTE_NEG=0x30,
    JSON_VALUE_BUF_TYPE_NUMBER_MULTIBYTE_NEG=0x40,
    JSON_VALUE_BUF_TYPE_STRING=0x50,
    JSON_VALUE_BUF_TYPE_ARRAY=0x60,
    JSON_VALUE_BUF_TYPE_OBJECT=0x70,
    JSON_VALUE_BUF_TYPE_SEPARATOR=0x80,
    JSON_VALUE_BUF_TYPE_SKIP=0xf0,
};

char *TOKEN_NAMES[] = {
    [JSON_TOKEN_NUMBER_INT] = "int",
    [JSON_TOKEN_NUMBER_FLOAT] = "float",
    [JSON_TOKEN_STRING] = "string",
    [JSON_TOKEN_LSQUARE] = "[",
    [JSON_TOKEN_RSQUARE] = "]",
    [JSON_TOKEN_LCURLY] = "{",
    [JSON_TOKEN_RCURLY] = "}",
    [JSON_TOKEN_COMMA] = ",",
    [JSON_TOKEN_COLON] = ":",
};
char *json_token_name(json_token t) {
    if (t <= JSON_TOKEN__START) return "error";
    return TOKEN_NAMES[t];
}

void json_lexer_new(json_lexer *l, u8 *buf, i64 size) {
    l->buf = buf;
    l->size = size;
    l->start = 0;
    l->end = 0;
}

static bool lex_is_more_input(json_lexer *l) { return l && l->end < l->size; }
static i64 lex_peek(json_lexer *l) { return l->buf[l->end]; }
static i64 lex_advance(json_lexer *l) {
    if (!lex_is_more_input(l)) return -1;
    return l->buf[l->end++];
}
static bool lex_match(json_lexer *l, i64 x) {
    if (lex_peek(l) == x) { lex_advance(l); return true; }
    return false;
}
static bool lex_match_pred(json_lexer *l, bool (*f)(i64)) {
    if (f(lex_peek(l))) { lex_advance(l); return true; }
    return false;
}

static bool is_digit(i64 x) { return x >= '0' && x <= '9'; }
static void consume_digits(json_lexer *l) { while (lex_match_pred(l, is_digit)); }
static json_token consume_float(json_lexer *l) {
    json_token ret = JSON_TOKEN_NUMBER_INT;
    consume_digits(l);
    if (lex_match(l, '.')) { consume_digits(l); ret = JSON_TOKEN_NUMBER_FLOAT; }
    if (lex_match(l, 'e') || lex_match(l, 'E')) {
        (void) (lex_match(l, '+') || lex_match(l, '-'));
        consume_digits(l);
        ret = JSON_TOKEN_NUMBER_FLOAT;
    }
    return ret;
}

json_token json_lex(json_lexer *l) {
start:
    l->start = l->end;
    i64 cur = lex_advance(l);
    switch (cur) {
    case '[': return JSON_TOKEN_LSQUARE;
    case ']': return JSON_TOKEN_RSQUARE;
    case '{': return JSON_TOKEN_LCURLY;
    case '}': return JSON_TOKEN_RCURLY;
    case ',': return JSON_TOKEN_COMMA;
    case ':': return JSON_TOKEN_COLON;
    case '"': {
        bool escaped = false;
        while ((cur = lex_advance(l))) {
            if (escaped) {
                escaped = false;
            } else {
                if (cur == '"') break;
                if (cur == '\\') escaped = true;
            }
        }
        if (cur < -1) return JSON_TOKEN__ERROR;
        return JSON_TOKEN_STRING;
    }
    case ' ': case '\t': case '\n': case '\r':
        l->buf[l->start] = JSON_VALUE_BUF_TYPE_SKIP;
        goto start;
    case '-':
        if (lex_match_pred(l, is_digit)) { return consume_float(l); }
        return JSON_TOKEN__ERROR;
    default:
        if (is_digit(cur)) { return consume_float(l); }
        return JSON_TOKEN__ERROR;
    }
}

void json_parser_new(json_parser *p, u8 *buf, i64 size) {
    json_lexer_new(&p->lex, buf, size);
    p->curstart = 0;
    p->curend = 0;
    p->next = json_lex(&p->lex);
}

static json_token parse_advance(json_parser *p) {
    if (p->next <= JSON_TOKEN__START) return p->next;
    json_token ret = p->next;
    p->curstart = p->lex.start;
    p->curend = p->lex.end;
    p->next = json_lex(&p->lex);
    return ret;
}
static json_token parse_peek(json_parser *p) { return p->next; }

static json_value val(i64 idx) { json_value ret; ret.idx = idx; return ret; }

static json_value install_string(json_parser *p, i64 start, i64 end) {
    p->lex.buf[start] = JSON_VALUE_BUF_TYPE_STRING;
    p->lex.buf[end - 1] = JSON_VALUE_BUF_TYPE_TERMINATOR;
    return val(start);
}
static json_value install_number_int(json_parser *p, i64 start, i64 end) {
    i64 x = 0;
    bool pos = true;
    i64 i = start;
    if (p->lex.buf[i] == '-') { pos = false; i += 1; }
    for (; i < end; ++i) {
        x *= 10; 
        x += p->lex.buf[i] - '0';
    }
    if (x < 16) {
        p->lex.buf[start] =
            (pos
                ? JSON_VALUE_BUF_TYPE_NUMBER_ONEBYTE_POS
                : JSON_VALUE_BUF_TYPE_NUMBER_ONEBYTE_NEG)
            | (u8) x;
    } else {
        i64 off = 1;
        for (; x > 0; off += 1, x >>= 8) {
            p->lex.buf[start + off] = (u8) x & 0xff;
        }
        p->lex.buf[start] = 
            (pos
                ? JSON_VALUE_BUF_TYPE_NUMBER_MULTIBYTE_POS
                : JSON_VALUE_BUF_TYPE_NUMBER_MULTIBYTE_NEG)
            | (u8) (off - 1);
    }
    return val(start);
}
static json_value install_array(json_parser *p, i64 start, i64 end) {
    p->lex.buf[start] = JSON_VALUE_BUF_TYPE_ARRAY;
    p->lex.buf[end - 1] = JSON_VALUE_BUF_TYPE_TERMINATOR;
    return val(start);
}

json_value_type json_value_get_type(json_parser *p, json_value v) {
    if (v.idx < 0 || v.idx >= p->lex.size) return JSON_VALUE_TYPE__ERROR;
    enum json_value_buf_type bt = p->lex.buf[v.idx] & 0xf0;
    switch (bt) {
    case JSON_VALUE_BUF_TYPE_NUMBER_ONEBYTE_POS:
    case JSON_VALUE_BUF_TYPE_NUMBER_MULTIBYTE_POS:
    case JSON_VALUE_BUF_TYPE_NUMBER_ONEBYTE_NEG:
    case JSON_VALUE_BUF_TYPE_NUMBER_MULTIBYTE_NEG:
        return JSON_VALUE_TYPE_NUMBER;
    case JSON_VALUE_BUF_TYPE_STRING: return JSON_VALUE_TYPE_STRING;
    case JSON_VALUE_BUF_TYPE_ARRAY: return JSON_VALUE_TYPE_ARRAY;
    case JSON_VALUE_BUF_TYPE_OBJECT: return JSON_VALUE_TYPE_OBJECT;
    default: return JSON_VALUE_TYPE__ERROR;
    }
}
char *json_value_as_string(json_parser *p, json_value v) {
    if (json_value_get_type(p, v) != JSON_VALUE_TYPE_STRING) return NULL;
    return (char *) &p->lex.buf[v.idx + 1];
}
i64 json_value_as_number(json_parser *p, json_value v) {
    enum json_value_buf_type bt = p->lex.buf[v.idx] & 0xf0;
    bool pos = true;
    bool multibyte = true;
    switch (bt) {
    case JSON_VALUE_BUF_TYPE_NUMBER_ONEBYTE_POS:
        multibyte = false;
        // fallthrough
    case JSON_VALUE_BUF_TYPE_NUMBER_MULTIBYTE_POS:
        break;
    case JSON_VALUE_BUF_TYPE_NUMBER_ONEBYTE_NEG:
        multibyte = false;
        // fallthrough
    case JSON_VALUE_BUF_TYPE_NUMBER_MULTIBYTE_NEG:
        pos = false;
        break;
    default: return 0;
    }
    if (multibyte) {
        i64 x = 0;
        i64 bs = p->lex.buf[v.idx] & 0x0f;
        for (i64 off = 0; off < bs; ++off) {
            x |= ((i64) p->lex.buf[v.idx + 1 + off]) << (off * 8);
        }
        return (i64) (pos ? 1 : -1) * x;
    } else {
        i64 x = p->lex.buf[v.idx] & 0x0f;
        return (i64) (pos ? 1 : -1) * x;
    }
}
json_array_iter json_value_as_array(json_parser *p, json_value v) {
    json_array_iter ret;
    if (json_value_get_type(p, v) != JSON_VALUE_TYPE_ARRAY) { ret.idx = -1; return ret; }
    ret.idx = v.idx + 1;
    return ret;
}
json_value json_array_iter_get(json_parser *p, json_array_iter i) {
    if (i.idx < 0 || i.idx >= p->lex.size) return val(-1);
    return val(i.idx);
}
bool json_array_iter_next(json_parser *p, json_array_iter *i) {
    i64 nesting = 0;
    bool in_string = false;
    while (i->idx < p->lex.size) {
        u8 b = p->lex.buf[i->idx];
        enum json_value_buf_type bt = b & 0xf0;
        i->idx += 1;
        if (in_string) {
            if (bt == JSON_VALUE_BUF_TYPE_TERMINATOR) in_string = false;
        } else {
            switch (bt) {
            case JSON_VALUE_BUF_TYPE_STRING: in_string = true; break;
            case JSON_VALUE_BUF_TYPE_NUMBER_MULTIBYTE_POS:
            case JSON_VALUE_BUF_TYPE_NUMBER_MULTIBYTE_NEG:
                i->idx += b & 0x0f;
                break;
            case JSON_VALUE_BUF_TYPE_ARRAY:
            case JSON_VALUE_BUF_TYPE_OBJECT:
                nesting += 1;
                break;
            case JSON_VALUE_BUF_TYPE_TERMINATOR:
                if (nesting > 0) nesting -= 1;
                else return false;
                break;
            case JSON_VALUE_BUF_TYPE_SEPARATOR:
                if (nesting == 0) {
                    while ((p->lex.buf[i->idx] & 0xf0) == JSON_VALUE_BUF_TYPE_SKIP) i->idx++;
                    return true;
                }
                break;
            default: break;
            }
        }
    }
    return false;
}

json_value json_parse(json_parser *p) {
    json_token cur = parse_advance(p);
    switch (cur) {
    case JSON_TOKEN_STRING: return install_string(p, p->curstart, p->curend);
    case JSON_TOKEN_NUMBER_INT: return install_number_int(p, p->curstart, p->curend);
    case JSON_TOKEN_LSQUARE: {
        i64 start = p->curstart;
        while (true) {
            json_parse(p);
            cur = parse_advance(p);
            if (cur == JSON_TOKEN_RSQUARE) break;
            else if (cur == JSON_TOKEN_COMMA) p->lex.buf[p->curstart] = JSON_VALUE_BUF_TYPE_SEPARATOR;
            else goto error;
        }
        return install_array(p, start, p->curend);
    }
    case JSON_TOKEN_LCURLY: {
        i64 start = p->curstart;
        while (true) {
            json_parse(p);
            cur = parse_advance(p);
            if (cur == JSON_TOKEN_RCURLY) break;
            else if (cur == JSON_TOKEN_COMMA || cur == JSON_TOKEN_COLON) p->lex.buf[p->curstart] = JSON_VALUE_BUF_TYPE_SEPARATOR;
            else goto error;
        }
        return install_array(p, start, p->curend);
    }
    default: goto error;
    }
error:
    return val(-1);
}