from stdint import *
import memhub
import string


# ============================================================
# JSON value type enum
# ============================================================
define JSON_NULL = 0
define JSON_BOOL = 1
define JSON_INT = 2
define JSON_FLOAT = 3
define JSON_STRING = 4
define JSON_ARRAY = 5
define JSON_OBJECT = 6


# ============================================================
# JsonValue - JSON value node.
# Allocated from mpool: JsonValue(pool) allocates from the pool.
# ============================================================
class JsonValue:
    vtype: int
    bool_val: int
    int_val: int64
    float_val: double
    str_val: str
    next: JsonValue*
    key: str
    child: JsonValue*
    child_count: sizet

    def __new__(self, pool: memhub.MemManager*):
        """Allocate and zero-initialize a JsonValue struct from mpool."""
        ptr: JsonValue* = pool.alloc(JsonValue.__sizeof__())
        if ptr:
            string.memset(ptr, 0, JsonValue.__sizeof__())
        return ptr

    def type(self) -> int:
        return self.vtype

    def is_null(self) -> bool:
        return self.vtype == JSON_NULL

    def is_bool(self) -> bool:
        return self.vtype == JSON_BOOL

    def is_int(self) -> bool:
        return self.vtype == JSON_INT

    def is_float(self) -> bool:
        return self.vtype == JSON_FLOAT

    def is_string(self) -> bool:
        return self.vtype == JSON_STRING

    def is_array(self) -> bool:
        return self.vtype == JSON_ARRAY

    def is_object(self) -> bool:
        return self.vtype == JSON_OBJECT

    def as_bool(self) -> bool:
        return self.bool_val != 0

    def as_int(self) -> int64:
        return self.int_val

    def as_float(self) -> double:
        return self.float_val

    def as_string(self) -> str:
        return self.str_val

    def __len__(self) -> sizet:
        return self.child_count

    def __getitem__(self, key: str) -> JsonValue*:
        """Python-style get: v["key"] returns an object field."""
        if self.vtype != JSON_OBJECT:
            return None
        if key == None:
            return None
        cur: JsonValue* = self.child
        while cur != None:
            if cur.key != None:
                if string.strcmp(cur.key, key) == 0:
                    return cur
            cur = cur.next
        return None

    def __setitem__(self, key: str, val: JsonValue*):
        """Python-style set: v["key"] = val."""
        if self.vtype != JSON_OBJECT or key == None or val == None:
            return
        cur: JsonValue* = self.child
        while cur != None:
            if cur.key != None:
                if string.strcmp(cur.key, key) == 0:
                    cur.vtype = val.vtype
                    cur.bool_val = val.bool_val
                    cur.int_val = val.int_val
                    cur.float_val = val.float_val
                    cur.str_val = val.str_val
                    cur.child = val.child
                    cur.child_count = val.child_count
                    return
            cur = cur.next
        # key absent: append a new child node
        klen: sizet = string.strlen(key)
        kbuf: str = self._pool_alloc(klen + 1)
        if kbuf != None:
            string.memcpy(kbuf, key, klen + 1)
            val.key = kbuf
        self._append_child(val)

    def get_item(self, index: sizet) -> JsonValue*:
        """Get an array element by index: v.get_item(0)."""
        if self.vtype != JSON_ARRAY and self.vtype != JSON_OBJECT:
            return None
        cur: JsonValue* = self.child
        i: sizet = 0
        while cur != None:
            if i == index:
                return cur
            cur = cur.next
            i += 1
        return None

    def _pool_alloc(self, size: sizet) -> str:
        """Internal: allocate from current object's mpool (reserved; currently returns None)."""
        return None

    def _append_child(self, child: JsonValue*):
        if self.child == None:
            self.child = child
        else:
            cur: JsonValue* = self.child
            while cur.next != None:
                cur = cur.next
            cur.next = child
        self.child_count += 1


# ============================================================
# Constructors: create each kind of JsonValue (allocated from mpool)
# ============================================================

def null(pool: memhub.MemManager*) -> JsonValue*:
    v: JsonValue* = JsonValue(pool)
    if v == None: return None
    v.vtype = JSON_NULL
    return v

def bool_val(pool: memhub.MemManager*, val: bool) -> JsonValue*:
    v: JsonValue* = JsonValue(pool)
    if v == None: return None
    v.vtype = JSON_BOOL
    v.bool_val = 1 if val else 0
    return v

def int_val(pool: memhub.MemManager*, val: int64) -> JsonValue*:
    v: JsonValue* = JsonValue(pool)
    if v == None: return None
    v.vtype = JSON_INT
    v.int_val = val
    return v

def float_val(pool: memhub.MemManager*, val: double) -> JsonValue*:
    v: JsonValue* = JsonValue(pool)
    if v == None: return None
    v.vtype = JSON_FLOAT
    v.float_val = val
    return v

def string_val(pool: memhub.MemManager*, val: str) -> JsonValue*:
    v: JsonValue* = JsonValue(pool)
    if v == None: return None
    v.vtype = JSON_STRING
    if val != None:
        slen: sizet = string.strlen(val)
        buf: str = pool.alloc(slen + 1)
        if buf != None:
            string.memcpy(buf, val, slen + 1)
            v.str_val = buf
    return v

def array(pool: memhub.MemManager*) -> JsonValue*:
    v: JsonValue* = JsonValue(pool)
    if v == None: return None
    v.vtype = JSON_ARRAY
    return v

def object(pool: memhub.MemManager*) -> JsonValue*:
    v: JsonValue* = JsonValue(pool)
    if v == None: return None
    v.vtype = JSON_OBJECT
    return v

# Append an element to an array
def array_append(pool: memhub.MemManager*, arr: JsonValue*, item: JsonValue*):
    if arr == None or item == None: return
    if arr.vtype != JSON_ARRAY: return
    arr._append_child(item)

# Set a key/value pair on an object
def object_set(pool: memhub.MemManager*, obj: JsonValue*, key: str, val: JsonValue*):
    if obj == None or key == None or val == None: return
    if obj.vtype != JSON_OBJECT: return
    cur: JsonValue* = obj.child
    while cur != None:
        if cur.key != None:
            if string.strcmp(cur.key, key) == 0:
                cur.vtype = val.vtype
                cur.bool_val = val.bool_val
                cur.int_val = val.int_val
                cur.float_val = val.float_val
                cur.str_val = val.str_val
                cur.child = val.child
                cur.child_count = val.child_count
                return
        cur = cur.next
    klen: sizet = string.strlen(key)
    kbuf: str = pool.alloc(klen + 1)
    if kbuf != None:
        string.memcpy(kbuf, key, klen + 1)
        val.key = kbuf
    obj._append_child(val)


# ============================================================
# Import submodules
# ============================================================
from .__parser import parse
from .__writer import write