include "std/containers.zr"

const json.OBJECT = 0
const json.ARRAY  = 1
const json.STRING = 2
const json.NUMBER = 3
const json.NULL   = 4
const json.BOOL   = 5

struct json.Value
    type: i64
    value: any

func json.Value.dump[v: json.Value] : str
    s := new str.Builder

    if v->type == json.OBJECT
        map := v->value as HashMap
        keys := map->keys()
        s->append_char('{')
        for i in 0..keys->size
            if i > 0
                s->append_char(',')
            s->append_char('"')
            s->append(keys->nth(i) as str)
            s->append_char('"')
            s->append_char(':')
            value, _ := map->get(keys->nth(i) as str)
            s->append((value as json.Value)->dump())
        s->append_char('}')
        keys->free()
    else if v->type == json.ARRAY
        arr := v->value as Array
        s->append_char('[')
        for i in 0..arr->size
            if i > 0
                s->append_char(',')
            s->append((arr->nth(i) as json.Value)->dump())
        s->append_char(']')
    else if v->type == json.STRING
        // TODO: escape sequences
        s->append_char('"')
        s->append(v->value as str)
        s->append_char('"')
    else if v->type == json.NUMBER
        s->append((v->value as i64)->to_str())
    else if v->type == json.NULL
        s->append("null")
    else if v->type == json.BOOL
        if v->value as bool
            s->append("true")
        else
            s->append("false")
    else
        panic("json.Value.dump: unknown value type")
    
    return s->build_and_free_buffer()

func json.Value.free[v: json.Value] : void
    if v->type == json.OBJECT
        map := v->value as HashMap
        keys := map->keys()
        for i in 0..keys->size
            value, _ := map->get(keys->nth(i) as str)
            (value as json.Value)->free()
        keys->free()
        
        map->free_with_keys()
    else if v->type == json.ARRAY
        arr := v->value as Array
        for i in 0..arr->size
            (arr->nth(i) as json.Value)->free()
        arr->free()
    else if v->type == json.STRING
        (v->value as str)->free()

    (v as ptr)->free()

func json.parse[s: str] : json.Value
    i := 0
    return json._parse_value(s, ^i)

func json._parse_value[s: str, i: ptr] : json.Value
    out := new* json.Value

    json._skip_whitespace(s, i)

    if s[mem.read64(i)] == 0
        panic("json.parse: unexpected EOF")
    
    if s[mem.read64(i)] == '{'
        mem.write64(i, mem.read64(i) + 1)
        json._skip_whitespace(s, i)

        out->type = json.OBJECT
        out->value = HashMap.new()
        
        if s[mem.read64(i)] != '}'
            while true
                json._skip_whitespace(s, i)
                key := json._parse_string(s, i)
                json._skip_whitespace(s, i)

                if s[mem.read64(i)] != ':'
                    panic("json.parse: expected ':'")
                mem.write64(i, mem.read64(i) + 1)
                json._skip_whitespace(s, i)

                value := json._parse_value(s, i)
                (out->value as HashMap)->insert(key, value)
                key->free()

                json._skip_whitespace(s, i)
                if s[mem.read64(i)] == ','
                    mem.write64(i, mem.read64(i) + 1)
                else
                    break
        
        json._skip_whitespace(s, i)
        if s[mem.read64(i)] != '}'
            panic("json.parse: expected '}'")
        mem.write64(i, mem.read64(i) + 1)
    else if s[mem.read64(i)] == '['
        mem.write64(i, mem.read64(i) + 1)
        json._skip_whitespace(s, i)

        out->type = json.ARRAY
        out->value = []
        
        while s[mem.read64(i)] != ']'
            json._skip_whitespace(s, i)

            (out->value as Array)->push(json._parse_value(s, i))
            json._skip_whitespace(s, i)

            if s[mem.read64(i)] == ','
                mem.write64(i, mem.read64(i) + 1)
            else
                break
        
        json._skip_whitespace(s, i)
        if s[mem.read64(i)] != ']'
            panic("json.parse: expected ']'")
        mem.write64(i, mem.read64(i) + 1)
    else if s[mem.read64(i)] == '"'
        out->type = json.STRING
        out->value = json._parse_string(s, i)
    else if s[mem.read64(i)]->is_digit() || s[mem.read64(i)] == '-'
        start := mem.read64(i)

        if s[mem.read64(i)] == '-'
            mem.write64(i, mem.read64(i) + 1)

        while s[mem.read64(i)]->is_digit()
            mem.write64(i, mem.read64(i) + 1)
        // TODO: float support once str.parse_f64 is a thing
        if s[mem.read64(i)] == '.'
            panic("json.parse: parsing floats not implemented yet")

        len := mem.read64(i) - start

        num_str := s->substr_n(start, len)
        out->type = json.NUMBER
        out->value = num_str->parse_i64()
        num_str->free()
    else if s[mem.read64(i)] == 'n' && s[mem.read64(i)+1] == 'u' && s[mem.read64(i)+2] == 'l' && s[mem.read64(i)+3] == 'l'
        out->type = json.NULL
        mem.write64(i, mem.read64(i) + 4)
    else if s[mem.read64(i)] == 't' && s[mem.read64(i)+1] == 'r' && s[mem.read64(i)+2] == 'u' && s[mem.read64(i)+3] == 'e'
        out->type = json.BOOL
        out->value = true
        mem.write64(i, mem.read64(i) + 4)
    else if s[mem.read64(i)] == 'f' && s[mem.read64(i)+1] == 'a' && s[mem.read64(i)+2] == 'l' && s[mem.read64(i)+3] == 's' && s[mem.read64(i)+4] == 'e'
        out->type = json.BOOL
        out->value = false
        mem.write64(i, mem.read64(i) + 5)
    else
        panic("json.parse: unexpected character")

    return out

func json._parse_string[s: str, i: ptr] : str
    if s[mem.read64(i)] != '"'
        panic("json.parse: expected '\"'")
    mem.write64(i, mem.read64(i) + 1)

    start := mem.read64(i)
    // TODO: escape sequences
    while s[mem.read64(i)] != '"'
        if s[mem.read64(i)] == 0
            panic("json.parse: unexpected EOF")
        mem.write64(i, mem.read64(i) + 1)

    len := mem.read64(i) - start
    mem.write64(i, mem.read64(i) + 1)
    return s->substr_n(start, len)

func json._skip_whitespace[s: str, i: ptr] : void
    while s[mem.read64(i)]->is_whitespace()
        mem.write64(i, mem.read64(i) + 1)