import t, c
from stdint import *
import memhub
import string
from stdio import snprintf
from . import (
    JsonValue, JSON_NULL, JSON_BOOL, JSON_INT, JSON_FLOAT,
    JSON_STRING, JSON_ARRAY, JSON_OBJECT,
)


def _write_char(buf: t.CChar | t.CPtr, pos: t.CSizeT, ch: t.CChar) -> t.CSizeT:
    buf[pos] = ch
    return pos + 1

def _write_str(buf: t.CChar | t.CPtr, pos: t.CSizeT, s: t.CChar | t.CPtr) -> t.CSizeT:
    if s == None: return pos
    slen: t.CSizeT = string.strlen(s)
    string.memcpy(t.CChar(t.CUInt64T(buf) + pos, t.CPtr), s, slen)
    return pos + slen

def _write_string_escaped(buf: t.CChar | t.CPtr, pos: t.CSizeT, s: t.CChar | t.CPtr) -> t.CSizeT:
    if s == None:
        pos = _write_str(buf, pos, "null")
        return pos
    pos = _write_char(buf, pos, '"')
    i: t.CSizeT = 0
    slen: t.CSizeT = string.strlen(s)
    while i < slen:
        ch: t.CChar = s[i]
        if ch == '"':
            pos = _write_str(buf, pos, "\\\"")
        elif ch == '\\':
            pos = _write_str(buf, pos, "\\\\")
        elif ch == '\n':
            pos = _write_str(buf, pos, "\\n")
        elif ch == '\r':
            pos = _write_str(buf, pos, "\\r")
        elif ch == '\t':
            pos = _write_str(buf, pos, "\\t")
        else:
            pos = _write_char(buf, pos, ch)
        i += 1
    pos = _write_char(buf, pos, '"')
    return pos


def _write_value(buf: t.CChar | t.CPtr, pos: t.CSizeT, val: JsonValue | t.CPtr, pool: memhub.MemManager | t.CPtr) -> t.CSizeT:
    if val == None:
        pos = _write_str(buf, pos, "null")
        return pos

    if val.vtype == JSON_NULL:
        pos = _write_str(buf, pos, "null")

    elif val.vtype == JSON_BOOL:
        if val.bool_val != 0:
            pos = _write_str(buf, pos, "true")
        else:
            pos = _write_str(buf, pos, "false")

    elif val.vtype == JSON_INT:
        tmp: t.CChar | t.CPtr = pool.alloc(32)
        if tmp != None:
            snprintf(tmp, 32, "%lld", val.int_val)
            pos = _write_str(buf, pos, tmp)

    elif val.vtype == JSON_FLOAT:
        tmp2: t.CChar | t.CPtr = pool.alloc(32)
        if tmp2 != None:
            snprintf(tmp2, 32, "%.6g", val.float_val)
            pos = _write_str(buf, pos, tmp2)

    elif val.vtype == JSON_STRING:
        pos = _write_string_escaped(buf, pos, val.str_val)

    elif val.vtype == JSON_ARRAY:
        if val.child == None:
            pos = _write_str(buf, pos, "[]")
            return pos
        pos = _write_char(buf, pos, '[')
        cur: JsonValue | t.CPtr = val.child
        first: bool = True
        while cur != None:
            if not first:
                pos = _write_char(buf, pos, ',')
            pos = _write_value(buf, pos, cur, pool)
            first = False
            cur = cur.next
        pos = _write_char(buf, pos, ']')

    elif val.vtype == JSON_OBJECT:
        if val.child == None:
            pos = _write_str(buf, pos, "{}")
            return pos
        pos = _write_char(buf, pos, '{')
        cur2: JsonValue | t.CPtr = val.child
        first2: bool = True
        while cur2 != None:
            if not first2:
                pos = _write_char(buf, pos, ',')
            pos = _write_string_escaped(buf, pos, cur2.key)
            pos = _write_char(buf, pos, ':')
            pos = _write_value(buf, pos, cur2, pool)
            first2 = False
            cur2 = cur2.next
        pos = _write_char(buf, pos, '}')

    return pos


def write(pool: memhub.MemManager | t.CPtr, val: JsonValue | t.CPtr, pretty: bool) -> t.CChar | t.CPtr:
    """将 JsonValue 序列化为 JSON 字符串"""
    buf: t.CChar | t.CPtr = pool.alloc(4096)
    if buf == None:
        return None
    end_pos: t.CSizeT = _write_value(buf, 0, val, pool)
    buf[end_pos] = '\0'
    return buf