import memhub
import string
from json 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,
    object_set,
)


# ============================================================
# tomllib - TOML parser producing a json.JsonValue tree.
#
# The JSON value model is structurally the same as TOML's (object/array/scalar), so TOML parses
# straight into JsonValue; writing it back out is json.write's job.
#
# Supported: [table] / [[array of tables]] headers, dotted keys (a.b.c = 1), inline tables,
# arrays (multi-line, with comments), basic + literal strings (single and triple quoted, with
# escapes and line-ending backslash), integers (decimal/0x/0o/0b, underscores, sign), floats
# (fraction/exponent), booleans, comments, blank lines.
#
# Deliberately rejected (TOML constructs with no json.JsonValue counterpart - never silently
# downgraded): date-time values (1979-05-27, 07:32:00, with Z/offset), and inf/nan.
#
# Return convention (same as json.parse): parse() returns a JSON_NULL node when the source is
# malformed or uses a rejected construct, and None only when the pool is exhausted. Callers must
# test `v == None or v.vtype == JSON_NULL` before walking the tree.
# ============================================================

# Maximum depth of a dotted key / table header path (stack-allocated segment array).
define TOML_MAX_KEY_SEGMENTS = 32
# Explicitly-defined table path registry capacity (bytes); overflow is a parse failure.
define TOML_DEFS_CAP = 65536


# TomlState - parser cursor plus the mpool all nodes are allocated from.
class TomlState:
    src: str
    pos: sizet = 0
    len: sizet = 0
    pool: memhub.MemManager*
    root: JsonValue* = None
    cur: JsonValue* = None
    cur_path: str = None
    defs: str = None
    defs_len: sizet = 0
    # 1 while the current section is an element of an [[array of tables]]: the registry below is
    # keyed on the dotted path only, so the same dotted key in two different elements of the same
    # array would otherwise look like a duplicate definition.
    in_array: int = 0

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


# ============================================================
# Character helpers
# ============================================================

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

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

def _is_space(ch: char) -> bool:
    return ch == ' ' or ch == '\t'

def _is_newline(ch: char) -> bool:
    return ch == '\n' or ch == '\r'

def _is_digit(ch: char) -> bool:
    return ch >= '0' and ch <= '9'

def _hex_val(ch: char) -> int:
    if _is_digit(ch): return int(ch) - 48
    if ch >= 'a' and ch <= 'f': return int(ch) - 87
    if ch >= 'A' and ch <= 'F': return int(ch) - 55
    return -1

def _is_bare_key_char(ch: char) -> bool:
    if ch >= 'a' and ch <= 'z': return True
    if ch >= 'A' and ch <= 'Z': return True
    if _is_digit(ch): return True
    return ch == '_' or ch == '-'

def _skip_spaces(s: TomlState*):
    while s.pos < s.len and _is_space(s.src[s.pos]): s.pos += 1

def _skip_to_line_end(s: TomlState*):
    while s.pos < s.len and not _is_newline(s.src[s.pos]): s.pos += 1

# Spaces, newlines and # comments: the filler between array items and between top-level lines.
def _skip_filler(s: TomlState*):
    while s.pos < s.len:
        ch: char = s.src[s.pos]
        if _is_space(ch) or _is_newline(ch): s.pos += 1
        elif ch == '#': _skip_to_line_end(s)
        else: break

# Trailing part of a completed statement: spaces, an optional comment, then the line must end.
def _skip_line_tail(s: TomlState*) -> int:
    _skip_spaces(s)
    if s.pos < s.len and s.src[s.pos] == '#': _skip_to_line_end(s)
    if s.pos >= s.len: return 0
    if _is_newline(s.src[s.pos]): s.pos += 1; return 0
    return -1

# Match a whole keyword (true/false) at the cursor without consuming it.
def _match_word(s: TomlState*, w: str) -> int:
    n: sizet = string.strlen(w)
    if s.pos + n > s.len: return 0
    if string.strncmp(s.src + s.pos, w, n) == 0: return 1
    return 0

# "..." / ''' open a multi-line string.
def _is_multiline_open(s: TomlState*, q: char) -> int:
    if s.pos + 2 < s.len and s.src[s.pos + 1] == q and s.src[s.pos + 2] == q: return 1
    return 0


# ============================================================
# Small allocation helpers
# ============================================================

def _copy_span(pool: memhub.MemManager*, src: str, start: sizet, n: sizet) -> str:
    buf: str = pool.alloc(n + 1)
    if buf == None: return None
    i: sizet = 0
    while i < n:
        buf[i] = src[start + i]
        i += 1
    buf[n] = '\0'
    return buf

# Encode one code point as UTF-8 (used by \uXXXX / \UXXXXXXXX); returns the new write index.
def _utf8_put(buf: str, di: sizet, cp: int) -> sizet:
    if cp < 0: cp = 0xFFFD
    if cp < 0x80:
        buf[di] = char(cp)
        return di + 1
    if cp < 0x800:
        buf[di] = char(0xC0 | (cp >> 6))
        buf[di + 1] = char(0x80 | (cp & 0x3F))
        return di + 2
    if cp < 0x10000:
        buf[di] = char(0xE0 | (cp >> 12))
        buf[di + 1] = char(0x80 | ((cp >> 6) & 0x3F))
        buf[di + 2] = char(0x80 | (cp & 0x3F))
        return di + 3
    buf[di] = char(0xF0 | (cp >> 18))
    buf[di + 1] = char(0x80 | ((cp >> 12) & 0x3F))
    buf[di + 2] = char(0x80 | ((cp >> 6) & 0x3F))
    buf[di + 3] = char(0x80 | (cp & 0x3F))
    return di + 4

# Read n hex digits that follow a \u / \U escape at s.pos; -1 on malformed input.
def _read_hex_escape(s: TomlState*, buf: str, di: sizet, n: sizet) -> int:
    cp: int = 0
    i: sizet = 0
    while i < n:
        p: sizet = s.pos + 1 + i
        if p >= s.len: return -1
        v: int = _hex_val(s.src[p])
        if v < 0: return -1
        cp = (cp << 4) | v
        i += 1
    s.pos = s.pos + n
    return int(_utf8_put(buf, di, cp))


# ============================================================
# Strings
# ============================================================

def _parse_string_escape(s: TomlState*, buf: str, di: sizet) -> int:
    """Decode the escape sequence after a backslash; returns the new write index, -1 on error."""
    if s.pos >= s.len: return -1
    esc: char = s.src[s.pos]
    if esc == 'b': buf[di] = char(8); return int(di) + 1
    if esc == 't': buf[di] = char(9); return int(di) + 1
    if esc == 'n': buf[di] = char(10); return int(di) + 1
    if esc == 'f': buf[di] = char(12); return int(di) + 1
    if esc == 'r': buf[di] = char(13); return int(di) + 1
    if esc == '"': buf[di] = '"'; return int(di) + 1
    if esc == '\\': buf[di] = '\\'; return int(di) + 1
    if esc == 'u': return _read_hex_escape(s, buf, di, 4)
    if esc == 'U': return _read_hex_escape(s, buf, di, 8)
    return -1

# Basic string "...": escapes are decoded. The output never exceeds the remaining source length
# (\uXXXX is 6 source chars for at most 4 output bytes), so a remaining+1 buffer always fits.
def _parse_basic_string(s: TomlState*) -> str:
    if _peek(s) != '"': return None
    _advance(s)
    start: sizet = s.pos
    buf: str = s.pool.alloc(s.len - start + 1)
    if buf == None: return None
    di: sizet = 0
    while s.pos < s.len:
        ch: char = s.src[s.pos]
        if ch == '"':
            s.pos += 1
            buf[di] = '\0'
            return buf
        if ch == '\\':
            s.pos += 1
            nd: int = _parse_string_escape(s, buf, di)
            if nd < 0: return None
            di = nd
            s.pos += 1
            continue
        if _is_newline(ch): return None
        buf[di] = ch
        di += 1
        s.pos += 1
    return None

# Literal string '...': taken verbatim, no escapes.
def _parse_literal_string(s: TomlState*) -> str:
    if _peek(s) != '\'': return None
    _advance(s)
    start: sizet = s.pos
    while s.pos < s.len and s.src[s.pos] != '\'':
        if _is_newline(s.src[s.pos]): return None
        s.pos += 1
    if s.pos >= s.len: return None
    n: sizet = s.pos - start
    s.pos += 1
    return _copy_span(s.pool, s.src, start, n)

# Multi-line basic string """...""": the first newline is trimmed; \ at end of line trims all
# following whitespace/newlines; escapes are decoded.
def _parse_multiline_basic(s: TomlState*) -> str:
    s.pos += 3
    if _peek(s) == '\r': s.pos += 1
    if _peek(s) == '\n': s.pos += 1
    start: sizet = s.pos
    buf: str = s.pool.alloc(s.len - start + 1)
    if buf == None: return None
    di: sizet = 0
    while s.pos < s.len:
        ch: char = s.src[s.pos]
        if ch == '"' and s.pos + 2 < s.len and s.src[s.pos + 1] == '"' and s.src[s.pos + 2] == '"':
            s.pos += 3
            buf[di] = '\0'
            return buf
        if ch == '\\':
            p: sizet = s.pos + 1
            while p < s.len and _is_space(s.src[p]): p += 1
            if p < s.len and _is_newline(s.src[p]):
                s.pos = p + 1
                while s.pos < s.len and (_is_space(s.src[s.pos]) or _is_newline(s.src[s.pos])): s.pos += 1
                continue
            s.pos += 1
            nd: int = _parse_string_escape(s, buf, di)
            if nd < 0: return None
            di = nd
            s.pos += 1
            continue
        buf[di] = ch
        di += 1
        s.pos += 1
    return None

# Multi-line literal string '''...''': verbatim except the trimmed leading newline.
def _parse_multiline_literal(s: TomlState*) -> str:
    s.pos += 3
    if _peek(s) == '\r': s.pos += 1
    if _peek(s) == '\n': s.pos += 1
    start: sizet = s.pos
    while s.pos < s.len:
        if s.src[s.pos] == '\'' and s.pos + 2 < s.len and s.src[s.pos + 1] == '\'' and s.src[s.pos + 2] == '\'':
            n: sizet = s.pos - start
            s.pos += 3
            return _copy_span(s.pool, s.src, start, n)
        s.pos += 1
    return None


# ============================================================
# Numbers
# ============================================================

# Classify a raw number token. Date-time and inf/nan are rejected (no JsonValue counterpart).
def _classify_number(pool: memhub.MemManager*, tok: str, n: sizet) -> JsonValue*:
    i: sizet = 0
    neg: int = 0
    if tok[0] == '-': neg = 1; i = 1
    elif tok[0] == '+': i = 1
    if i >= n: return None
    # No JSON counterpart: date/time (':' or 'Z'), and inf/nan.
    if tok[i] == 'i' or tok[i] == 'n': return None
    j: sizet = i
    while j < n:
        ch: char = tok[j]
        if ch == ':' or ch == 'T' or ch == 'Z': return None
        # A '-'/'+' may only follow an exponent marker: "1e-3" is a float, "2024-05-27" is a date
        # (rejected: json.JsonValue has no date-time kind).
        if ch == '-' or ch == '+':
            if tok[j - 1] != 'e' and tok[j - 1] != 'E': return None
        j += 1
    # 0x / 0o / 0b integers
    if i + 1 < n and tok[i] == '0':
        base_ch: char = tok[i + 1]
        if base_ch == 'x' or base_ch == 'o' or base_ch == 'b':
            base: int = 16
            if base_ch == 'o': base = 8
            elif base_ch == 'b': base = 2
            acc: int64 = 0
            k: sizet = i + 2
            digits: sizet = 0
            while k < n:
                ch2: char = tok[k]
                if ch2 != '_':
                    v: int = _hex_val(ch2)
                    if v < 0 or v >= base: return None
                    acc = acc * base + v
                    digits += 1
                k += 1
            if digits == 0: return None
            if neg != 0: acc = -acc
            return int_val(pool, acc)
    # decimal integer or float: strip '_', keep [0-9.eE+-]
    clean: str = pool.alloc(n - i + 1)
    if clean == None: return None
    di: sizet = 0
    is_float: int = 0
    dots: int = 0
    exps: int = 0
    k2: sizet = i
    while k2 < n:
        ch3: char = tok[k2]
        if ch3 == '_': k2 += 1; continue
        if ch3 == '.':
            dots += 1
            if dots > 1: return None                   # at most one fraction separator
            is_float = 1
        elif ch3 == 'e' or ch3 == 'E':
            exps += 1
            if exps > 1: return None                   # at most one exponent
            is_float = 1
        elif not _is_digit(ch3) and ch3 != '+' and ch3 != '-': return None
        clean[di] = ch3
        di += 1
        k2 += 1
    if di == 0: return None
    clean[di] = '\0'
    # Leading zeros: only a lone "0" may start the integer part, so "01" and "00.5" are rejected
    # while "0.5" and "0e0" (integer part "0") are fine.
    ip: sizet = 0
    while ip < di and _is_digit(clean[ip]): ip += 1
    if ip > 1 and clean[0] == '0': return None
    if is_float != 0:
        tailc: char = clean[di - 1]
        if tailc == '.' or tailc == 'e' or tailc == 'E' or tailc == '+' or tailc == '-': return None
        fv: double = string.atof(clean)
        if neg != 0: fv = 0.0 - fv
        return float_val(pool, fv)
    iv: int64 = string.atoll(clean)
    if neg != 0: iv = -iv
    return int_val(pool, iv)

# Number token: run to the next delimiter, then classify.
def _parse_number_value(s: TomlState*) -> JsonValue*:
    start: sizet = s.pos
    while s.pos < s.len:
        ch: char = s.src[s.pos]
        if _is_space(ch) or _is_newline(ch): break
        if ch == ',' or ch == ']' or ch == '}' or ch == '#': break
        s.pos += 1
    n: sizet = s.pos - start
    if n == 0: return None
    tok: str = _copy_span(s.pool, s.src, start, n)
    if tok == None: return None
    return _classify_number(s.pool, tok, n)


# ============================================================
# Containers and values
# ============================================================

def _parse_array(s: TomlState*) -> JsonValue*:
    _advance(s)                                        # '['
    arr: JsonValue* = array(s.pool)
    if arr == None: return None
    while True:
        _skip_filler(s)
        if s.pos >= s.len: return None
        if _peek(s) == ']': _advance(s); return arr
        item: JsonValue* = _parse_value(s)
        if item == None: return None
        arr._append_child(item)
        _skip_filler(s)
        if s.pos < s.len and _peek(s) == ',':
            _advance(s)
            continue
        if s.pos < s.len and _peek(s) == ']': _advance(s); return arr
        return None

def _parse_inline_table(s: TomlState*) -> JsonValue*:
    _advance(s)                                        # '{'
    obj: JsonValue* = object(s.pool)
    if obj == None: return None
    _skip_spaces(s)
    if s.pos < s.len and _peek(s) == '}': _advance(s); return obj
    segs: [TOML_MAX_KEY_SEGMENTS]str
    while True:
        n: int = _parse_key_path(s, segs, TOML_MAX_KEY_SEGMENTS)
        if n <= 0: return None
        _skip_spaces(s)
        if _peek(s) != '=': return None
        _advance(s)
        val: JsonValue* = _parse_value(s)
        if val == None: return None
        tbl: JsonValue* = _dotted_target(s, obj, segs, int(n), 0)
        if tbl == None: return None
        last: str = segs[n - 1]
        if tbl[last] != None: return None
        object_set(s.pool, tbl, last, val)
        if tbl[last] == None: return None
        _skip_spaces(s)
        if s.pos < s.len and _peek(s) == '}': _advance(s); return obj
        if s.pos < s.len and _peek(s) == ',': _advance(s); _skip_spaces(s); continue
        return None

def _parse_value(s: TomlState*) -> JsonValue*:
    _skip_spaces(s)
    if s.pos >= s.len: return None
    ch: char = _peek(s)
    if ch == '"':
        sv: str = None
        if _is_multiline_open(s, '"') != 0: sv = _parse_multiline_basic(s)
        else: sv = _parse_basic_string(s)
        if sv == None: return None
        return string_val(s.pool, sv)
    if ch == '\'':
        sv2: str = None
        if _is_multiline_open(s, '\'') != 0: sv2 = _parse_multiline_literal(s)
        else: sv2 = _parse_literal_string(s)
        if sv2 == None: return None
        return string_val(s.pool, sv2)
    if ch == '[': return _parse_array(s)
    if ch == '{': return _parse_inline_table(s)
    if ch == 't':
        if _match_word(s, "true") != 0: s.pos += 4; return bool_val(s.pool, True)
        return None
    if ch == 'f':
        if _match_word(s, "false") != 0: s.pos += 5; return bool_val(s.pool, False)
        return None
    if ch == '+' or ch == '-' or _is_digit(ch): return _parse_number_value(s)
    return None


# ============================================================
# Keys and table paths
# ============================================================

def _parse_key_segment(s: TomlState*) -> str:
    _skip_spaces(s)
    ch: char = _peek(s)
    if ch == '"':
        if _is_multiline_open(s, '"') != 0: return None
        return _parse_basic_string(s)
    if ch == '\'':
        if _is_multiline_open(s, '\'') != 0: return None
        return _parse_literal_string(s)
    if not _is_bare_key_char(ch): return None
    start: sizet = s.pos
    while s.pos < s.len and _is_bare_key_char(s.src[s.pos]): s.pos += 1
    return _copy_span(s.pool, s.src, start, s.pos - start)

# Dotted key path into segs; returns the segment count, -1 on error.
def _parse_key_path(s: TomlState*, segs: [TOML_MAX_KEY_SEGMENTS]str, cap: int) -> int:
    n: int = 0
    while True:
        if n >= cap: return -1
        seg: str = _parse_key_segment(s)
        if seg == None: return -1
        segs[n] = seg
        n += 1
        _skip_spaces(s)
        if _peek(s) != '.': return n
        _advance(s)

# Join segs[0..n) into a dotted path, prefixed by `prefix` when it is a non-empty path.
def _join_path(pool: memhub.MemManager*, prefix: str, segs: [TOML_MAX_KEY_SEGMENTS]str, n: int) -> str:
    total: sizet = 1
    if prefix != None: total += string.strlen(prefix) + 1
    i: int = 0
    while i < n:
        total += string.strlen(segs[i]) + 1
        i += 1
    buf: str = pool.alloc(total)
    if buf == None: return None
    di: sizet = 0
    if prefix != None and prefix[0] != '\0':
        plen: sizet = string.strlen(prefix)
        string.memcpy(buf, prefix, plen)
        di = plen
        buf[di] = '.'
        di += 1
    i = 0
    while i < n:
        if i > 0: buf[di] = '.'; di += 1
        slen: sizet = string.strlen(segs[i])
        string.memcpy(buf + di, segs[i], slen)
        di += slen
        i += 1
    buf[di] = '\0'
    return buf

# ============================================================
# Explicitly-defined table registry (TOML forbids defining a table twice)
# ============================================================

def _defs_init(s: TomlState*) -> int:
    s.defs = s.pool.alloc(TOML_DEFS_CAP)
    if s.defs == None: return -1
    s.defs[0] = '\0'
    s.defs_len = 0
    return 0

def _defs_contains(s: TomlState*, path: str) -> int:
    if s.defs == None or path == None: return 0
    plen: sizet = string.strlen(path)
    pat: str = s.pool.alloc(plen + 3)
    if pat == None: return 0
    pat[0] = '\n'
    string.memcpy(pat + 1, path, plen)
    pat[plen + 1] = '\n'
    pat[plen + 2] = '\0'
    if string.strstr(s.defs, pat) != None: return 1
    return 0

def _defs_add(s: TomlState*, path: str) -> int:
    if s.defs == None or path == None: return -1
    plen: sizet = string.strlen(path)
    if s.defs_len + plen + 3 > TOML_DEFS_CAP: return -1
    s.defs[s.defs_len] = '\n'
    s.defs_len += 1
    string.memcpy(s.defs + s.defs_len, path, plen)
    s.defs_len += plen
    s.defs[s.defs_len] = '\n'
    s.defs_len += 1
    s.defs[s.defs_len] = '\0'
    return 0

def _mark_defined(s: TomlState*, path: str) -> int:
    if _defs_contains(s, path) != 0: return -1
    return _defs_add(s, path)


# ============================================================
# Table navigation
# ============================================================

# Child table under `parent`: existing object, or a freshly created one. None on a name clash.
def _child_table(pool: memhub.MemManager*, parent: JsonValue*, key: str) -> JsonValue*:
    if parent == None or key == None: return None
    exist: JsonValue* = parent[key]
    if exist != None:
        if exist.vtype == JSON_OBJECT: return exist
        return None
    fresh: JsonValue* = object(pool)
    if fresh == None: return None
    object_set(pool, parent, key, fresh)
    if parent[key] == None: return None                # key buffer allocation failed
    return fresh

# Walk segs[0..n-1) creating intermediate tables. With mark != 0 a **newly created** intermediate
# is registered as explicitly defined (a dotted key defines its table, so a later [table] header for
# the same path is an error), while a reused table is left alone: `a.b = 1` then `a.c = 2` is legal.
# Inline-table keys pass mark = 0 because they are outside the table registry.
def _dotted_target(s: TomlState*, base: JsonValue*, segs: [TOML_MAX_KEY_SEGMENTS]str, n: int, mark: int) -> JsonValue*:
    tbl: JsonValue* = base
    i: int = 0
    while i < n - 1:
        existed: int = 0
        raw: JsonValue* = tbl[segs[i]]
        if raw != None and raw.vtype == JSON_OBJECT: existed = 1
        nxt: JsonValue* = _child_table(s.pool, tbl, segs[i])
        if nxt == None: return None
        if mark != 0 and existed == 0 and s.in_array == 0:
            sub: str = _join_path(s.pool, s.cur_path, segs, i + 1)
            if sub == None: return None
            if _mark_defined(s, sub) != 0: return None
        tbl = nxt
        i += 1
    return tbl

# [[key]]: append a new object element to the array `key` of `parent`, creating the array first.
def _array_element_new(pool: memhub.MemManager*, parent: JsonValue*, key: str) -> JsonValue*:
    if parent == None or key == None: return None
    arr: JsonValue* = parent[key]
    if arr == None:
        arr = array(pool)
        if arr == None: return None
        object_set(pool, parent, key, arr)
        arr = parent[key]
        if arr == None: return None
    elif arr.vtype != JSON_ARRAY:
        return None
    elem: JsonValue* = object(pool)
    if elem == None: return None
    arr._append_child(elem)
    return elem


# ============================================================
# Statements
# ============================================================

# [a.b.c] - absolute path from the root; the parent chain must already be tables.
# [[a.b.c]] - the last segment is an array of tables: append a new element and switch to it.
def _parse_table_header(s: TomlState*) -> int:
    _advance(s)                                        # '['
    is_arr: int = 0
    if _peek(s) == '[': is_arr = 1; _advance(s)
    segs: [TOML_MAX_KEY_SEGMENTS]str
    n: int = _parse_key_path(s, segs, TOML_MAX_KEY_SEGMENTS)
    if n <= 0: return -1
    _skip_spaces(s)
    if _peek(s) != ']': return -1
    _advance(s)
    if is_arr != 0:
        if _peek(s) != ']': return -1
        _advance(s)
    if _skip_line_tail(s) != 0: return -1
    parent: JsonValue* = s.root
    i: int = 0
    while i < n - 1:
        parent = _child_table(s.pool, parent, segs[i])
        if parent == None: return -1
        i += 1
    last: str = segs[n - 1]
    if is_arr != 0:
        elem: JsonValue* = _array_element_new(s.pool, parent, last)
        if elem == None: return -1
        tail: str = _join_path(s.pool, "", segs, n)
        if tail == None: return -1
        s.cur = elem
        s.cur_path = tail
        s.in_array = 1
        return 0
    path: str = _join_path(s.pool, "", segs, n)
    if path == None: return -1
    if _mark_defined(s, path) != 0: return -1
    tbl: JsonValue* = _child_table(s.pool, parent, last)
    if tbl == None: return -1
    s.cur = tbl
    s.cur_path = path
    s.in_array = 0
    return 0

# key = value, relative to the current table; dotted keys create intermediate tables.
def _parse_keyval(s: TomlState*) -> int:
    segs: [TOML_MAX_KEY_SEGMENTS]str
    n: int = _parse_key_path(s, segs, TOML_MAX_KEY_SEGMENTS)
    if n <= 0: return -1
    _skip_spaces(s)
    if _peek(s) != '=': return -1
    _advance(s)
    val: JsonValue* = _parse_value(s)
    if val == None: return -1
    tbl: JsonValue* = _dotted_target(s, s.cur, segs, int(n), 1)
    if tbl == None: return -1
    last: str = segs[n - 1]
    if tbl[last] != None: return -1                   # duplicate key
    object_set(s.pool, tbl, last, val)
    if tbl[last] == None: return -1
    return _skip_line_tail(s)


# ============================================================
# Entry point
# ============================================================

def parse(pool: memhub.MemManager*, src: str) -> JsonValue*:
    """Parse a TOML document into a JsonValue tree (JSON_NULL node on failure, None only if the pool is exhausted)."""
    if src == None: return null(pool)
    s: TomlState* = TomlState(src, pool)
    if _defs_init(s) != 0: return null(pool)
    s.root = object(pool)
    if s.root == None: return null(pool)
    s.cur = s.root
    s.cur_path = ""
    failed: int = 0
    while True:
        _skip_filler(s)
        if s.pos >= s.len: break
        if _peek(s) == '[':
            if _parse_table_header(s) != 0: failed = 1; break
        else:
            if _parse_keyval(s) != 0: failed = 1; break
    if failed != 0: return null(pool)
    return s.root