import t, c
from stdint import *
import memhub
import string
from . import (
    JsonValue, JSON_NULL, JSON_BOOL, JSON_INT, JSON_FLOAT,
    JSON_STRING, JSON_ARRAY, JSON_OBJECT,
    null, bool_val, int_val, float_val, string_val, array, object,
)


class ParseState:
    src: t.CChar | t.CPtr
    pos: t.CSizeT
    len: t.CSizeT
    pool: memhub.MemManager | t.CPtr

    def __init__(self, src: t.CChar | t.CPtr, pool: memhub.MemManager | t.CPtr):
        self.src = src
        self.pos = 0
        self.len = string.strlen(src) if src != None else 0
        self.pool = pool


def _peek(s: ParseState | t.CPtr) -> t.CChar:
    if s.pos >= s.len:
        return '\0'
    return s.src[s.pos]

def _advance(s: ParseState | t.CPtr) -> t.CChar:
    if s.pos >= s.len:
        return '\0'
    ch: t.CChar = s.src[s.pos]
    s.pos += 1
    return ch

def _skip_whitespace(s: ParseState | t.CPtr):
    while s.pos < s.len:
        ch: t.CChar = s.src[s.pos]
        if ch == ' ' or ch == '\t' or ch == '\n' or ch == '\r':
            s.pos += 1
        else:
            break

def _expect(s: ParseState | t.CPtr, ch: t.CChar) -> bool:
    _skip_whitespace(s)
    if s.pos >= s.len:
        return False
    if s.src[s.pos] != ch:
        return False
    s.pos += 1
    return True


def _parse_string(s: ParseState | t.CPtr) -> t.CChar | t.CPtr:
    if _peek(s) != '"':
        return None
    _advance(s)

    start: t.CSizeT = s.pos
    slen: t.CSizeT = 0
    while s.pos < s.len:
        ch: t.CChar = s.src[s.pos]
        if ch == '"':
            break
        if ch == '\\':
            s.pos += 1
            slen += 1
        slen += 1
        s.pos += 1

    buf: t.CChar | t.CPtr = s.pool.alloc(slen + 1)
    if buf == None:
        return None

    s.pos = start
    di: t.CSizeT = 0
    while s.pos < s.len:
        ch: t.CChar = s.src[s.pos]
        if ch == '"':
            s.pos += 1
            break
        if ch == '\\':
            s.pos += 1
            if s.pos < s.len:
                esc: t.CChar = s.src[s.pos]
                if esc == 'n':
                    buf[di] = '\n'
                elif esc == 't':
                    buf[di] = '\t'
                elif esc == 'r':
                    buf[di] = '\r'
                elif esc == '\\':
                    buf[di] = '\\'
                elif esc == '"':
                    buf[di] = '"'
                elif esc == '/':
                    buf[di] = '/'
                else:
                    buf[di] = esc
                di += 1
                s.pos += 1
            continue
        buf[di] = ch
        di += 1
        s.pos += 1

    buf[di] = '\0'
    return buf


def _parse_number(s: ParseState | t.CPtr) -> JsonValue | t.CPtr:
    _skip_whitespace(s)
    is_float: bool = False
    start: t.CSizeT = s.pos

    if s.pos < s.len and s.src[s.pos] == '-':
        s.pos += 1

    while s.pos < s.len:
        ch: t.CChar = s.src[s.pos]
        if ch >= '0' and ch <= '9':
            s.pos += 1
        else:
            break

    if s.pos < s.len and s.src[s.pos] == '.':
        is_float = True
        s.pos += 1
        while s.pos < s.len:
            ch: t.CChar = s.src[s.pos]
            if ch >= '0' and ch <= '9':
                s.pos += 1
            else:
                break

    if s.pos < s.len and (s.src[s.pos] == 'e' or s.src[s.pos] == 'E'):
        is_float = True
        s.pos += 1
        if s.pos < s.len and (s.src[s.pos] == '+' or s.src[s.pos] == '-'):
            s.pos += 1
        while s.pos < s.len:
            ch: t.CChar = s.src[s.pos]
            if ch >= '0' and ch <= '9':
                s.pos += 1
            else:
                break

    num_len: t.CSizeT = s.pos - start
    buf: t.CChar | t.CPtr = s.pool.alloc(num_len + 1)
    if buf == None:
        return null(s.pool)

    i: t.CSizeT = 0
    while i < num_len:
        buf[i] = s.src[start + i]
        i += 1
    buf[num_len] = '\0'

    if is_float:
        val: t.CDouble = string.atof(buf)
        return float_val(s.pool, val)
    else:
        ival: t.CInt64T = string.atoll(buf)
        return int_val(s.pool, ival)


def _parse_value(s: ParseState | t.CPtr) -> JsonValue | t.CPtr:
    _skip_whitespace(s)
    if s.pos >= s.len:
        return null(s.pool)

    ch: t.CChar = _peek(s)

    if ch == '"':
        sv: t.CChar | t.CPtr = _parse_string(s)
        return string_val(s.pool, sv)

    if ch == '{':
        return _parse_object(s)

    if ch == '[':
        return _parse_array(s)

    if ch == '-' or (ch >= '0' and ch <= '9'):
        return _parse_number(s)

    if ch == 't':
        if s.pos + 3 < s.len:
            if s.src[s.pos + 1] == 'r' and s.src[s.pos + 2] == 'u' and s.src[s.pos + 3] == 'e':
                s.pos += 4
                return bool_val(s.pool, True)
        return null(s.pool)

    if ch == 'f':
        if s.pos + 4 < s.len:
            if s.src[s.pos + 1] == 'a' and s.src[s.pos + 2] == 'l' and s.src[s.pos + 3] == 's' and s.src[s.pos + 4] == 'e':
                s.pos += 5
                return bool_val(s.pool, False)
        return null(s.pool)

    if ch == 'n':
        if s.pos + 3 < s.len:
            if s.src[s.pos + 1] == 'u' and s.src[s.pos + 2] == 'l' and s.src[s.pos + 3] == 'l':
                s.pos += 4
                return null(s.pool)
        return null(s.pool)

    return null(s.pool)


def _parse_array(s: ParseState | t.CPtr) -> JsonValue | t.CPtr:
    if not _expect(s, '['):
        return null(s.pool)

    arr: JsonValue | t.CPtr = array(s.pool)

    _skip_whitespace(s)
    if s.pos < s.len and _peek(s) == ']':
        _advance(s)
        return arr

    item: JsonValue | t.CPtr = _parse_value(s)
    arr._append_child(item)

    while True:
        _skip_whitespace(s)
        if s.pos >= s.len:
            break
        if _peek(s) != ',':
            break
        _advance(s)
        item = _parse_value(s)
        arr._append_child(item)

    _expect(s, ']')
    return arr


def _parse_object(s: ParseState | t.CPtr) -> JsonValue | t.CPtr:
    if not _expect(s, '{'):
        return null(s.pool)

    obj: JsonValue | t.CPtr = object(s.pool)

    _skip_whitespace(s)
    if s.pos < s.len and _peek(s) == '}':
        _advance(s)
        return obj

    _skip_whitespace(s)
    key: t.CChar | t.CPtr = _parse_string(s)
    _expect(s, ':')
    val: JsonValue | t.CPtr = _parse_value(s)
    if key != None and val != None:
        klen: t.CSizeT = string.strlen(key)
        kbuf: t.CChar | t.CPtr = s.pool.alloc(klen + 1)
        if kbuf != None:
            string.memcpy(kbuf, key, klen + 1)
            val.key = kbuf
        obj._append_child(val)

    while True:
        _skip_whitespace(s)
        if s.pos >= s.len:
            break
        if _peek(s) != ',':
            break
        _advance(s)
        _skip_whitespace(s)
        key = _parse_string(s)
        _expect(s, ':')
        val = _parse_value(s)
        if key != None and val != None:
            klen2: t.CSizeT = string.strlen(key)
            kbuf2: t.CChar | t.CPtr = s.pool.alloc(klen2 + 1)
            if kbuf2 != None:
                string.memcpy(kbuf2, key, klen2 + 1)
                val.key = kbuf2
            obj._append_child(val)

    _expect(s, '}')
    return obj


def parse(pool: memhub.MemManager | t.CPtr, src: t.CChar | t.CPtr) -> JsonValue | t.CPtr:
    """解析 JSON 字符串，返回 JsonValue 树"""
    if src == None:
        return null(pool)
    s: ParseState | t.CPtr = ParseState(src, pool)
    return _parse_value(s)