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: str
    pos: sizet = 0
    len: sizet = 0
    pool: memhub.MemManager*

    def __init__(self, src: str, pool: memhub.MemManager*):
        self.src = src
        self.len = string.strlen(src) if src != None else 0
        self.pool = pool


def _peek(s: ParseState*) -> char:
    if s.pos >= s.len: return '\0'
    return s.src[s.pos]

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

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

def _expect(s: ParseState*, ch: char) -> 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*) -> str:
    if _peek(s) != '"': return None
    _advance(s)

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

    buf: str = s.pool.alloc(slen + 1)
    if buf == None:
        return None

    s.pos = start
    di: sizet = 0
    while s.pos < s.len:
        ch: char = s.src[s.pos]
        if ch == '"':
            s.pos += 1
            break
        if ch == '\\':
            s.pos += 1
            if s.pos < s.len:
                esc: char = 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*) -> JsonValue*:
    _skip_whitespace(s)
    is_float: bool = False
    start: sizet = s.pos

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

    while s.pos < s.len:
        ch: char = 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: char = 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: char = s.src[s.pos]
            if ch >= '0' and ch <= '9':
                s.pos += 1
            else:
                break

    num_len: sizet = s.pos - start
    buf: str = s.pool.alloc(num_len + 1)
    if buf == None:
        return null(s.pool)

    for i in range(int(num_len)):
        buf[i] = s.src[start + i]
    buf[num_len] = '\0'

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


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

    ch: char = _peek(s)

    if ch == '"':
        sv: str = _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*) -> JsonValue*:
    if not _expect(s, '['):
        return null(s.pool)

    arr: JsonValue* = array(s.pool)
    # Pool exhausted => array() returns None; continuing to arr._append_child(...) would deref
    # a null pointer, so return None here (json.parse may return None; callers must check).
    if arr == None:
        return null(s.pool)

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

    item: JsonValue* = _parse_value(s)
    # item == None iff allocation failed somewhere inside _parse_value (syntax errors yield
    # a valid JSON_NULL node; only pool exhaustion yields None) => return None, append nothing.
    if item == None:
        return null(s.pool)
    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)
        if item == None:
            return null(s.pool)
        arr._append_child(item)

    _expect(s, ']')
    return arr


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

    obj: JsonValue* = object(s.pool)
    # Same as _parse_array: obj is None on pool exhaustion; do not call _append_child on None.
    if obj == None:
        return null(s.pool)

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

    _skip_whitespace(s)
    key: str = _parse_string(s)
    _expect(s, ':')
    val: JsonValue* = _parse_value(s)
    if key != None and val != None:
        klen: sizet = string.strlen(key)
        kbuf: str = 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: sizet = string.strlen(key)
            kbuf2: str = 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*, src: str) -> JsonValue*:
    """Parse a JSON string, returning a JsonValue tree."""
    if src == None:
        return null(pool)
    s: ParseState* = ParseState(src, pool)
    return _parse_value(s)