import "./core.zc"
import "./vec.zc"
import "./map.zc"
import "./string.zc"
import "./result.zc"
import "./option.zc"
import "./mem.zc"

@derive(Eq)
enum JsonType {
    JSON_NULL,
    JSON_BOOL,
    JSON_NUMBER,
    JSON_STRING,
    JSON_ARRAY,
    JSON_OBJECT
}

struct JsonValue {
    kind: JsonType;
    string_val: char*;
    number_val: f64;
    bool_val: bool;
    array_val: Vec<JsonValue*>*;
    object_val: Map<JsonValue*>*;
}

alias JsonValuePtr = JsonValue*;

fn _json_skip_ws(data: char*, pos: usize*) {
    while (data[*pos] == ' ' || data[*pos] == '\t' ||
           data[*pos] == '\n' || data[*pos] == '\r') {
        *pos = *pos + 1;
    }
}

fn _json_parse_string(data: char*, pos: usize*) -> char* {
    if (data[*pos] != '"') return NULL;
    *pos = *pos + 1;
    let v = Vec<char>::new();
    while (data[*pos] != 0 && data[*pos] != '"') {
        if (data[*pos] == '\\' && data[*pos + 1] != 0) {
            *pos = *pos + 1;
            let c = data[*pos];
            if (c == 'n') v.push('\n');
            else if (c == 't') v.push('\t');
            else if (c == 'r') v.push('\r');
            else if (c == '"') v.push('"');
            else if (c == '\\') v.push('\\');
            else if (c == '/') v.push('/');
            else if (c == 'b') v.push('\b');
            else if (c == 'f') v.push('\f');
            else v.push(c);
        } else {
            v.push(data[*pos]);
        }
        *pos = *pos + 1;
    }
    if (data[*pos] == '"') {
        *pos = *pos + 1;
    }
    v.push(0);
    let result = v.data;
    v.forget();
    return result;
}

fn _json_parse_value(data: char*, pos: usize*) -> JsonValue*;
fn _json_parse_array(data: char*, pos: usize*) -> JsonValue*;
fn _json_parse_object(data: char*, pos: usize*) -> JsonValue*;

fn _json_make_null() -> JsonValue* {
    let v: JsonValue* = alloc<JsonValue>();
    *v = JsonValue { kind: JsonType::JSON_NULL(), string_val: NULL, number_val: 0.0, bool_val: false, array_val: NULL, object_val: NULL };
    return v;
}

fn _json_make_bool(b: bool) -> JsonValue* {
    let v: JsonValue* = alloc<JsonValue>();
    *v = JsonValue { kind: JsonType::JSON_BOOL(), string_val: NULL, number_val: 0.0, bool_val: b, array_val: NULL, object_val: NULL };
    return v;
}

fn _json_make_string(s: char*) -> JsonValue* {
    let v: JsonValue* = alloc<JsonValue>();
    *v = JsonValue { kind: JsonType::JSON_STRING(), string_val: s, number_val: 0.0, bool_val: false, array_val: NULL, object_val: NULL };
    return v;
}

fn _json_make_number(n: f64) -> JsonValue* {
    let v: JsonValue* = alloc<JsonValue>();
    *v = JsonValue { kind: JsonType::JSON_NUMBER(), string_val: NULL, number_val: n, bool_val: false, array_val: NULL, object_val: NULL };
    return v;
}

fn _json_make_array(v: Vec<JsonValue*>*) -> JsonValue* {
    let arr: JsonValue* = alloc<JsonValue>();
    *arr = JsonValue { kind: JsonType::JSON_ARRAY(), string_val: NULL, number_val: 0.0, bool_val: false, array_val: v, object_val: NULL };
    return arr;
}

fn _json_make_object(m: Map<JsonValue*>*) -> JsonValue* {
    let obj: JsonValue* = alloc<JsonValue>();
    *obj = JsonValue { kind: JsonType::JSON_OBJECT(), string_val: NULL, number_val: 0.0, bool_val: false, array_val: NULL, object_val: m };
    return obj;
}

fn _json_parse_value(data: char*, pos: usize*) -> JsonValue* {
    _json_skip_ws(data, pos);
    if (data[*pos] == 0) return NULL;

    if (strncmp(data + *pos, "null", 4) == 0) {
        *pos = *pos + 4;
        return _json_make_null();
    }
    if (strncmp(data + *pos, "true", 4) == 0) {
        *pos = *pos + 4;
        return _json_make_bool(true);
    }
    if (strncmp(data + *pos, "false", 5) == 0) {
        *pos = *pos + 5;
        return _json_make_bool(false);
    }
    if (data[*pos] == '"') {
        let s = _json_parse_string(data, pos);
        if (s == NULL) return NULL;
        return _json_make_string(s);
    }
    if (data[*pos] == '-' || (data[*pos] >= '0' && data[*pos] <= '9')) {
        let end: char*;
        let num = strtod(data + *pos, &end);
        if (end == data + *pos) return NULL;
        *pos = (usize)(end - data);
        return _json_make_number(num);
    }
    if (data[*pos] == '[') {
        return _json_parse_array(data, pos);
    }
    if (data[*pos] == '{') {
        return _json_parse_object(data, pos);
    }
    return NULL;
}

fn _json_parse_array(data: char*, pos: usize*) -> JsonValue* {
    if (data[*pos] != '[') return NULL;
    *pos = *pos + 1;
    let v: Vec<JsonValue*>* = alloc<Vec<JsonValue*>>();
    *v = Vec__JsonValuePtr::new();

    _json_skip_ws(data, pos);
    if (data[*pos] == ']') {
        *pos = *pos + 1;
        return _json_make_array(v);
    }

    while (true) {
        _json_skip_ws(data, pos);
        let val = _json_parse_value(data, pos);
        if (val == NULL) return NULL;
        Vec__JsonValuePtr__push(v, val);
        _json_skip_ws(data, pos);
        if (data[*pos] == ']') {
            *pos = *pos + 1;
            return _json_make_array(v);
        }
        if (data[*pos] != ',') return NULL;
        *pos = *pos + 1;
    }
}

fn _json_parse_object(data: char*, pos: usize*) -> JsonValue* {
    if (data[*pos] != '{') return NULL;
    *pos = *pos + 1;
    let m: Map<JsonValue*>* = alloc<Map<JsonValue*>>();
    *m = Map__JsonValuePtr::new();

    _json_skip_ws(data, pos);
    if (data[*pos] == '}') {
        *pos = *pos + 1;
        return _json_make_object(m);
    }

    while (true) {
        _json_skip_ws(data, pos);
        let key = _json_parse_string(data, pos);
        if (key == NULL) return NULL;
        _json_skip_ws(data, pos);
        if (data[*pos] != ':') {
            dealloc<char>(key);
            return NULL;
        }
        *pos = *pos + 1;
        _json_skip_ws(data, pos);
        let val = _json_parse_value(data, pos);
        if (val == NULL) {
            dealloc<char>(key);
            return NULL;
        }
        Map__JsonValuePtr__put(m, key, val);
        dealloc<char>(key);
        _json_skip_ws(data, pos);
        if (data[*pos] == '}') {
            *pos = *pos + 1;
            return _json_make_object(m);
        }
        if (data[*pos] != ',') return NULL;
        *pos = *pos + 1;
    }
}

fn _json_do_parse(json: char*) -> JsonValue* {
    let pos: usize = 0;
    return _json_parse_value(json, &pos);
}

impl JsonValue {
    /// Creates a JSON null value.
    fn null() -> JsonValue {
        return JsonValue { kind: JsonType::JSON_NULL(), string_val: NULL, number_val: 0.0, bool_val: false, array_val: NULL, object_val: NULL };
    }
    
    /// Creates a JSON boolean.
    fn bool(b: bool) -> JsonValue {
        return JsonValue { kind: JsonType::JSON_BOOL(), string_val: NULL, number_val: 0.0, bool_val: b, array_val: NULL, object_val: NULL };
    }
    
    /// Creates a JSON number.
    fn number(n: f64) -> JsonValue {
        return JsonValue { kind: JsonType::JSON_NUMBER(), string_val: NULL, number_val: n, bool_val: false, array_val: NULL, object_val: NULL };
    }
    
    /// Creates a JSON string.
    fn string(s: char*) -> JsonValue {
        let len = strlen(s);
        let dup: char* = alloc_n<char>(len + 1);
        if (dup != NULL) {
            memcpy(dup, s, len + 1);
        }
        return JsonValue { kind: JsonType::JSON_STRING(), string_val: dup, number_val: 0.0, bool_val: false, array_val: NULL, object_val: NULL };
    }
    
    /// Creates a JSON array.
    fn array() -> JsonValue {
        let v: Vec<JsonValue*>* = alloc<Vec<JsonValue*>>();
        *v = Vec__JsonValuePtr::new();
        return JsonValue { kind: JsonType::JSON_ARRAY(), string_val: NULL, number_val: 0.0, bool_val: false, array_val: v, object_val: NULL };
    }
    
    /// Creates a JSON object.
    fn object() -> JsonValue {
        let m: Map<JsonValue*>* = alloc<Map<JsonValue*>>();
        *m = Map__JsonValuePtr::new();
        return JsonValue { kind: JsonType::JSON_OBJECT(), string_val: NULL, number_val: 0.0, bool_val: false, array_val: NULL, object_val: m };
    }

    // Heap-allocated factory methods (returns pointers)

    /// Creates a heap-allocated JSON null value.
    fn null_ptr() -> JsonValue* {
        let p: JsonValue* = alloc<JsonValue>();
        *p = JsonValue::null();
        return p;
    }

    /// Creates a heap-allocated JSON boolean.
    fn bool_ptr(b: bool) -> JsonValue* {
        let p: JsonValue* = alloc<JsonValue>();
        *p = JsonValue::bool(b);
        return p;
    }

    /// Creates a heap-allocated JSON number.
    fn number_ptr(n: f64) -> JsonValue* {
        let p: JsonValue* = alloc<JsonValue>();
        *p = JsonValue::number(n);
        return p;
    }

    /// Creates a heap-allocated JSON string.
    fn string_ptr(s: char*) -> JsonValue* {
        let p: JsonValue* = alloc<JsonValue>();
        *p = JsonValue::string(s);
        return p;
    }

    /// Creates a heap-allocated JSON array.
    fn array_ptr() -> JsonValue* {
        let p: JsonValue* = alloc<JsonValue>();
        *p = JsonValue::array();
        return p;
    }

    /// Creates a heap-allocated JSON object.
    fn object_ptr() -> JsonValue* {
        let p: JsonValue* = alloc<JsonValue>();
        *p = JsonValue::object();
        return p;
    }

    /// Appends a value to the array.
    fn push(self, val: JsonValue) {
        if (self.kind.tag != JsonType::JSON_ARRAY().tag) return;
        let p: JsonValue* = alloc<JsonValue>();
        *p = val;
        self.array_val.push(p);
    }
    
    /// Sets an object field.
    fn set(self, key: char*, val: JsonValue) {
        if (self.kind.tag != JsonType::JSON_OBJECT().tag) return;
        let p: JsonValue* = alloc<JsonValue>();
        *p = val;
        self.object_val.put(key, p);
    }
    
    /// Parses a JSON string into a JsonValue.
    fn parse(json: char*) -> Result<JsonValue*> {
        let result: JsonValue* = _json_do_parse(json);
        if (result != NULL) {
            return Result<JsonValue*>::Ok(result);
        }
        return Result<JsonValue*>::Err("JSON parse error");
    }

    /// Parses a JSON string into an owned JsonValue.
    fn parse_val(json: char*) -> Result<JsonValue> {
        let result: JsonValue* = _json_do_parse(json);
        if (result != NULL) {
            let val: JsonValue = *result;
            dealloc<JsonValue>(result);
            return Result<JsonValue>::Ok(val);
        }
        return Result<JsonValue>::Err("JSON parse error");
    }

    // Type checking helpers

    /// Returns whether the value is null.
    fn is_null(self) -> bool {
        return self.kind.tag == JsonType::JSON_NULL().tag;
    }

    /// Returns whether the value is a boolean.
    fn is_bool(self) -> bool {
        return self.kind.tag == JsonType::JSON_BOOL().tag;
    }

    /// Returns whether the value is a number.
    fn is_number(self) -> bool {
        return self.kind.tag == JsonType::JSON_NUMBER().tag;
    }

    /// Returns whether the value is a string.
    fn is_string(self) -> bool {
        return self.kind.tag == JsonType::JSON_STRING().tag;
    }

    /// Returns whether the value is an array.
    fn is_array(self) -> bool {
        return self.kind.tag == JsonType::JSON_ARRAY().tag;
    }

    /// Returns whether the value is an object.
    fn is_object(self) -> bool {
        return self.kind.tag == JsonType::JSON_OBJECT().tag;
    }

    // Direct value extractors

    /// Returns the string value when it is a string.
    fn as_string(self) -> Option<char*> {
        if self.kind.tag == JsonType::JSON_STRING().tag {
            return Option<char*>::Some(self.string_val);
        }
        return Option<char*>::None();
    }

    /// Returns the integer value when it is a number.
    fn as_int(self) -> Option<int> {
        if self.kind.tag == JsonType::JSON_NUMBER().tag {
            return Option<int>::Some((int)self.number_val);
        }
        return Option<int>::None();
    }

    /// Returns the floating-point value when it is a number.
    fn as_float(self) -> Option<f64> {
        if self.kind.tag == JsonType::JSON_NUMBER().tag {
            return Option<f64>::Some(self.number_val);
        }
        return Option<f64>::None();
    }

    /// Returns the boolean value when it is a boolean.
    fn as_bool(self) -> Option<bool> {
        if self.kind.tag == JsonType::JSON_BOOL().tag {
            return Option<bool>::Some(self.bool_val);
        }
        return Option<bool>::None();
    }

    // Object key accessors

    /// Returns the object field value for `key`.
    fn get(self, key: char*) -> Option<JsonValue*> {
        if self.kind.tag != JsonType::JSON_OBJECT().tag {
            return Option<JsonValue*>::None();
        }
        if Map<JsonValue*>::contains(self.object_val, key) {
            return Map<JsonValue*>::get(self.object_val, key);
        }
        return Option<JsonValue*>::None();
    }

    /// Returns the string field for `key`, or `None`.
    fn get_string(self, key: char*) -> Option<char*> {
        let opt = self.get(key);
        if opt.is_none() {
            return Option<char*>::None();
        }
        let val = opt.unwrap();
        return (*val).as_string();
    }

    /// Returns the integer field for `key`, or `None`.
    fn get_int(self, key: char*) -> Option<int> {
        let opt = self.get(key);
        if opt.is_none() {
            return Option<int>::None();
        }
        let val = opt.unwrap();
        return (*val).as_int();
    }

    /// Returns the floating-point field for `key`, or `None`.
    fn get_float(self, key: char*) -> Option<f64> {
        let opt = self.get(key);
        if opt.is_none() {
            return Option<f64>::None();
        }
        let val = opt.unwrap();
        return (*val).as_float();
    }

    /// Returns the boolean field for `key`, or `None`.
    fn get_bool(self, key: char*) -> Option<bool> {
        let opt = self.get(key);
        if opt.is_none() {
            return Option<bool>::None();
        }
        let val = opt.unwrap();
        return (*val).as_bool();
    }

    /// Returns the object field for `key`, or `None`.
    fn get_object(self, key: char*) -> Option<JsonValue*> {
        let opt = self.get(key);
        if opt.is_none() {
            return Option<JsonValue*>::None();
        }
        let val = opt.unwrap();
        if (*val).kind.tag == JsonType::JSON_OBJECT().tag {
            return Option<JsonValue*>::Some(val);
        }
        return Option<JsonValue*>::None();
    }

    /// Returns the array field for `key`, or `None`.
    fn get_array(self, key: char*) -> Option<JsonValue*> {
        let opt = self.get(key);
        if opt.is_none() {
            return Option<JsonValue*>::None();
        }
        let val = opt.unwrap();
        if (*val).kind.tag == JsonType::JSON_ARRAY().tag {
            return Option<JsonValue*>::Some(val);
        }
        return Option<JsonValue*>::None();
    }

    // Array accessors

    /// Returns the array element at `index`.
    fn at(self, index: usize) -> Option<JsonValue*> {
        if self.kind.tag != JsonType::JSON_ARRAY().tag {
            return Option<JsonValue*>::None();
        }
        if index >= self.array_val.length() {
            return Option<JsonValue*>::None();
        }
        return Option<JsonValue*>::Some(self.array_val.get(index));
    }

    /// Returns the array/object size.
    fn length(self) -> usize {
        if self.kind.tag == JsonType::JSON_ARRAY().tag {
            return self.array_val.length();
        }
        if self.kind.tag == JsonType::JSON_OBJECT().tag {
            return self.object_val.length();
        }
        return 0;
    }

    // Deprecated alias; use `length`.
    /// Returns the array/object size (deprecated; use `length`).
    fn len(self) -> usize {
        return self.length();
    }

    /// Recursively frees the value and its children.
    fn free(self) {
        if (self.kind.tag == JsonType::JSON_STRING().tag) dealloc<char>(self.string_val);
        if (self.kind.tag == JsonType::JSON_ARRAY().tag) {
             let v: Vec<JsonValue*>* = self.array_val;
             for (let i: usize = 0; i < v.length(); i = i + 1) {
                 let item: JsonValue* = v.get(i);
                 item.free();
                 dealloc<JsonValue>(item);
             }
             v.clear();
             dealloc<JsonValue*>(v.data);
             dealloc<Vec<JsonValue*>>(v);
        }
        if (self.kind.tag == JsonType::JSON_OBJECT().tag) {
             let m: Map<JsonValue*>* = self.object_val;
             for (let i: usize = 0; i < m.capacity(); i = i + 1) {
                 if (m.is_slot_occupied(i)) {
                     let child: JsonValue* = m.val_at(i);
                     child.free();
                     dealloc<JsonValue>(child);
                 }
             }
             m.free();
             dealloc<Map<JsonValue*>>(m);
        }
    }
}

impl Drop for JsonValue {
    fn drop(self) {
        self.free();
    }
}

impl JsonValue {
    /// Returns a serialized JSON string (caller frees).
    fn to_string(self) -> char* {
        let s = String::new("");
        self.stringify(&s);
        let out = s.c_str();
        s.forget();
        return out;
    }

    /// Serializes the value into the given String buffer.
    fn stringify(self, buf: String*) {
        if (self.kind.tag == JsonType::JSON_NULL().tag) {
            buf.append_c_ptr("null");
        } else if (self.kind.tag == JsonType::JSON_BOOL().tag) {
            if (self.bool_val) { buf.append_c_ptr("true"); } else { buf.append_c_ptr("false"); }
        } else if (self.kind.tag == JsonType::JSON_NUMBER().tag) {
            let tmp: char[64];
            sprintf((char*)tmp, "%.15g", self.number_val);
            buf.append_c_ptr((char*)tmp);
        } else if (self.kind.tag == JsonType::JSON_STRING().tag) {
            buf.append_c_ptr("\"");
            let p = self.string_val;
            let len = strlen(p);
            for (let i = 0; i < len; i = i + 1) {
                let c = p[i];
                if (c == '"') buf.append_c_ptr("\\\"");
                else if (c == '\\') buf.append_c_ptr("\\\\");
                else if (c == '\n') buf.append_c_ptr("\\n");
                else if (c == '\t') buf.append_c_ptr("\\t");
                else if (c == '\r') buf.append_c_ptr("\\r");
                else if (c == '\b') buf.append_c_ptr("\\b");
                else if (c == '\f') buf.append_c_ptr("\\f");
                else {
                    let tmp: char[2]; tmp[0] = c; tmp[1] = 0;
                    buf.append_c_ptr((char*)tmp);
                }
            }
            buf.append_c_ptr("\"");
        } else if (self.kind.tag == JsonType::JSON_ARRAY().tag) {
            buf.append_c_ptr("[");
            let v = self.array_val;
            for (let i: usize = 0; i < v.length(); i = i + 1) {
                if (i > 0) buf.append_c_ptr(",");
                let item = v.get(i);
                (*item).stringify(buf);
            }
            buf.append_c_ptr("]");
        } else if (self.kind.tag == JsonType::JSON_OBJECT().tag) {
            buf.append_c_ptr("{{");
            let m = self.object_val;
            let first = true;
            for (let i: usize = 0; i < m.capacity(); i = i + 1) {
                if (m.is_slot_occupied(i)) {
                    if (!first) buf.append_c_ptr(",");
                    first = false;
                    let key = m.key_at(i);
                    buf.append_c_ptr("\"");
                    buf.append_c_ptr(key);
                    buf.append_c_ptr("\":");
                    let val = m.val_at(i);
                    val.stringify(buf);
                }
            }
            buf.append_c_ptr("}");
        }
    }
}