"""
SPDX-License-Identifier: BSL-1.0

Copyright (c) 2025 Juan Jin <juanjin.dev@gmail.com>
"""

from __future__ import annotations
import enum
from dataclasses import dataclass, field
from typing import Any, List, Optional, Iterable

# ---------- Token ----------
class JdaughterTokType(enum.Enum):
    LBRACE = "{"
    RBRACE = "}"
    LBRACKET = "["
    RBRACKET = "]"
    COLON = ":"
    COMMA = ","
    STRING = "STRING"
    NUMBER = "NUMBER"
    TRUE = "TRUE"
    FALSE = "FALSE"
    NULL = "NULL"
    NAN = "NAN"
    INLINE_COMMENT = "INLINE_COMMENT"
    BLOCK_COMMENT = "BLOCK_COMMENT"
    EOF = "EOF"

@dataclass
class JdaughterToken:
    type: "JdaughterTokType"
    lexeme: str
    pos: int
    line: int  # 1-based line number

class JdaughterLexerError(Exception):
    ...

class JdaughterParseError(Exception):
    ...

# ---------- Lexer ----------
class JdaughterLexer:
    def __init__(self, src: str):
        self.s = src
        self.n = len(src)
        self.i = 0
        self.line = 1  # 1-based
        self._last_char = ""

    def _p(self, k=0):
        j = self.i + k
        return self.s[j] if j < self.n else ""

    def _adv(self):
        c = self._p()
        self.i += 1
        if c == "\n":
            self.line += 1
        self._last_char = c
        return c

    def _ws(self):
        while self.i < self.n and self.s[self.i] in " \t\r\n":
            self._adv()

    def _string(self) -> "JdaughterToken":
        start = self.i
        start_line = self.line
        self._adv()  # opening "
        out = []
        while True:
            if self.i >= self.n:
                raise JdaughterLexerError("unterminated string")

            ch = self._adv()

            if ch in ("\n", "\r"):
                raise JdaughterLexerError("newline in string")

            if ch == '"':
                break

            if ch == "\\":
                if self.i >= self.n:
                    raise JdaughterLexerError("bad escape")
                esc = self._adv()
                mp = {
                    '"': '"',
                    "\\": "\\",
                    "/": "/",
                    "b": "\b",
                    "f": "\f",
                    "n": "\n",
                    "r": "\r",
                    "t": "\t",
                }
                if esc in mp:
                    out.append(mp[esc])
                elif esc == "u":
                    he = self.s[self.i : self.i + 4]
                    if len(he) < 4 or any(
                        c not in "0123456789abcdefABCDEF" for c in he
                    ):
                        raise JdaughterLexerError("bad \\u")
                    out.append(chr(int(he, 16)))
                    self.i += 4
                else:
                    raise JdaughterLexerError(f"bad \\{esc}")
            else:
                out.append(ch)

        return JdaughterToken(
            JdaughterTokType.STRING, "".join(out), start, start_line
        )

    def _number(self) -> "JdaughterToken":
        start = self.i
        start_line = self.line
        s: list[str] = []

        # optional sign
        if self._p() == "-":
            s.append(self._adv())

        digits = "0123456789"

        # ----- integer part -----
        first = self._p()
        if first == "0":
            s.append(self._adv())
            nxt = self._p()
            if nxt and nxt in digits:
                raise JdaughterLexerError("leading zero")
        else:
            if first not in digits:
                raise JdaughterLexerError("number")
            s.append(self._adv())
            while True:
                c = self._p()
                if not c or c not in digits:
                    break
                s.append(self._adv())

        # ----- fraction part -----
        c = self._p()
        if c == ".":
            s.append(self._adv())
            c = self._p()
            if not c or c not in digits:
                raise JdaughterLexerError("frac")
            s.append(self._adv())
            while True:
                c = self._p()
                if not c or c not in digits:
                    break
                s.append(self._adv())

        # ----- exponent part -----
        c = self._p()
        if c and c in "eE":
            s.append(self._adv())  # consume e/E

            c = self._p()
            if c and c in "+-":
                s.append(self._adv())
                c = self._p()

            if not c or c not in digits:
                raise JdaughterLexerError("exp")

            s.append(self._adv())
            while True:
                c = self._p()
                if not c or c not in digits:
                    break
                s.append(self._adv())

        return JdaughterToken(JdaughterTokType.NUMBER, "".join(s), start, start_line)



    def _inline(self) -> "JdaughterToken":
        start = self.i
        start_line = self.line
        self._adv()
        self._adv()  # //
        buf = []
        while self.i < self.n and self._p() not in "\n\r":
            buf.append(self._adv())
        return JdaughterToken(
            JdaughterTokType.INLINE_COMMENT, "".join(buf), start, start_line
        )

    def _block(self) -> "JdaughterToken":
        start = self.i
        start_line = self.line
        self._adv()
        self._adv()  # /*
        buf = []
        while True:
            if self.i >= self.n:
                raise JdaughterLexerError("unterminated block comment")
            if self._p() == "*" and self._p(1) == "/":
                self._adv()
                self._adv()
                break
            buf.append(self._adv())
        return JdaughterToken(
            JdaughterTokType.BLOCK_COMMENT, "".join(buf), start, start_line
        )

    def tokens(self) -> Iterable["JdaughterToken"]:
        while True:
            self._ws()
            if self.i >= self.n:
                break
            c = self._p()
            line = self.line
            if c == "{":
                self._adv()
                yield JdaughterToken(
                    JdaughterTokType.LBRACE, "{", self.i - 1, line
                )
            elif c == "}":
                self._adv()
                yield JdaughterToken(
                    JdaughterTokType.RBRACE, "}", self.i - 1, line
                )
            elif c == "[":
                self._adv()
                yield JdaughterToken(
                    JdaughterTokType.LBRACKET, "[", self.i - 1, line
                )
            elif c == "]":
                self._adv()
                yield JdaughterToken(
                    JdaughterTokType.RBRACKET, "]", self.i - 1, line
                )
            elif c == ":":
                self._adv()
                yield JdaughterToken(
                    JdaughterTokType.COLON, ":", self.i - 1, line
                )
            elif c == ",":
                self._adv()
                yield JdaughterToken(
                    JdaughterTokType.COMMA, ",", self.i - 1, line
                )
            elif c == '"':
                yield self._string()
            elif c == "-" or c.isdigit():
                yield self._number()
            elif c == "/":
                if self._p(1) == "/":
                    yield self._inline()
                elif self._p(1) == "*":
                    yield self._block()
                else:
                    raise JdaughterLexerError("bad /")
            else:
                # idents: true / false / null / NaN
                start = self.i
                start_line = self.line
                ident = []
                while self._p().isalpha() or self._p() == "_":
                    ident.append(self._adv())
                ident_str = "".join(ident)
                tt = {
                    "true": JdaughterTokType.TRUE,
                    "false": JdaughterTokType.FALSE,
                    "null": JdaughterTokType.NULL,
                    "NaN": JdaughterTokType.NAN,
                }.get(ident_str)
                if tt:
                    yield JdaughterToken(tt, ident_str, start, start_line)
                else:
                    raise JdaughterLexerError(f"unexpected {ident_str!r}")
        yield JdaughterToken(JdaughterTokType.EOF, "", self.i, self.line)

# ---------- AST ----------
class JdaughterType(enum.Enum):
    DOCUMENT = 0
    OBJECT = 1
    ARRAY = 2
    PAIR = 3
    KEY = 4
    STRING = 5
    NUMBER = 6
    BOOLEAN = 7
    NULL = 8
    NAN = 9
    INLINE_COMMENT = 10
    BLOCK_COMMENT = 11

@dataclass
class JdaughterNode:
    type: "JdaughterType"
    value: Any = None
    children: List[Any] = field(default_factory=list)

    # ---- DOCUMENT unwrap ----
    def _unwrap_root(self) -> "JdaughterNode":
        if self.type == JdaughterType.DOCUMENT:
            for ch in self.children:
                if isinstance(ch, JdaughterNode) and ch.type not in (
                    JdaughterType.INLINE_COMMENT,
                    JdaughterType.BLOCK_COMMENT,
                ):
                    return ch
            raise ValueError("DOCUMENT has no root value")
        return self

    # ---- operator overloading ----
    def __getitem__(self, key):
        if isinstance(key, tuple):
            node = self
            for k in key:
                node = node[k]
            return node

        node = self._unwrap_root()
        T = JdaughterType
        if node.type == T.OBJECT:
            if not isinstance(key, str):
                raise TypeError("object indexing expects str key")
            for e in node.children:  # JdaughterObjEntry
                if (
                    e.pair
                    and e.pair.children
                    and e.pair.children[0].type == T.KEY
                    and e.pair.children[0].value == key
                ):
                    return e.pair.children[1]
            raise KeyError(key)

        if node.type == T.ARRAY:
            if not isinstance(key, int):
                raise TypeError("array indexing expects int index")
            n = len(node.children)
            if key < 0:
                key += n
            if key < 0 or key >= n:
                raise IndexError(key)
            entry = node.children[key]
            if entry.value is None:
                raise IndexError("array entry has no value")
            return entry.value

        raise TypeError(f"__getitem__ not supported for {node.type}")

    def __setitem__(self, key, value):
        if isinstance(key, tuple) and len(key) > 0:
            *head, last = key
            target = self
            for k in head:
                target = target[k]
            target[last] = value
            return

        node = self._unwrap_root()
        T = JdaughterType
        wrapped = _wrap(value)

        if node.type == T.OBJECT:
            if not isinstance(key, str):
                raise TypeError("object indexing expects str key")
            for e in node.children:
                if (
                    e.pair
                    and e.pair.children
                    and e.pair.children[0].type == T.KEY
                    and e.pair.children[0].value == key
                ):
                    e.pair.children[1] = wrapped
                    return
            pair = JdaughterNode(
                T.PAIR,
                None,
                [JdaughterNode(T.KEY, key), wrapped],
            )
            node.children.append(JdaughterObjEntry(pair=pair))
            return

        if node.type == T.ARRAY:
            if not isinstance(key, int):
                raise TypeError("array indexing expects int index")
            n = len(node.children)
            idx = key if key >= 0 else n + key
            if idx == n:
                node.children.append(JdaughterArrEntry(value=wrapped))
                return
            if idx < 0 or idx >= n:
                raise IndexError(key)
            node.children[idx].value = wrapped
            return

        raise TypeError(f"__setitem__ not supported for {node.type}")

    def __delitem__(self, key):
        node = self._unwrap_root()
        T = JdaughterType

        if isinstance(key, tuple) and len(key) > 0:
            *head, last = key
            target = node
            for k in head:
                target = target[k]
            target.__delitem__(last)
            return

        if node.type == T.OBJECT:
            if not isinstance(key, str):
                raise TypeError("object key must be str")
            new_entries = []
            removed = False
            for e in node.children:
                if (
                    e.pair
                    and e.pair.children
                    and e.pair.children[0].value == key
                ):
                    removed = True
                    continue
                new_entries.append(e)
            if not removed:
                raise KeyError(key)
            node.children = new_entries
            return

        if node.type == T.ARRAY:
            if not isinstance(key, int):
                raise TypeError("array index must be int")
            n = len(node.children)
            idx = key if key >= 0 else n + key
            if idx < 0 or idx >= n:
                raise IndexError(key)
            node.children.pop(idx)
            return

        raise TypeError(f"__delitem__ not supported for {node.type}")

    def __contains__(self, k):
        node = self._unwrap_root()
        T = JdaughterType
        if node.type == T.OBJECT and isinstance(k, str):
            return any(
                e.pair
                and e.pair.children
                and e.pair.children[0].value == k
                for e in node.children
            )
        if node.type == T.ARRAY and isinstance(k, int):
            return -len(node.children) <= k < len(node.children)
        return False

    # dict-like helpers
    def keys(self):
        node = self._unwrap_root()
        if node.type != JdaughterType.OBJECT:
            raise TypeError("keys() only for OBJECT")
        return [e.pair.children[0].value for e in node.children if e.pair]

    def values(self):
        node = self._unwrap_root()
        if node.type != JdaughterType.OBJECT:
            raise TypeError("values() only for OBJECT")
        return [e.pair.children[1] for e in node.children if e.pair]

    def items(self):
        node = self._unwrap_root()
        if node.type != JdaughterType.OBJECT:
            raise TypeError("items() only for OBJECT")
        return [
            (e.pair.children[0].value, e.pair.children[1])
            for e in node.children
            if e.pair
        ]

    def get(self, key, default=None):
        try:
            return self[key]
        except KeyError:
            return default

    def setdefault(self, key: str, default_value=None):
        node = self._unwrap_root()
        if node.type != JdaughterType.OBJECT:
            raise TypeError("setdefault() only for OBJECT")
        try:
            return self[key]
        except KeyError:
            self[key] = default_value
            return self[key]

    def update(self, mapping: dict):
        node = self._unwrap_root()
        if node.type != JdaughterType.OBJECT:
            raise TypeError("update() only for OBJECT")
        for k, v in mapping.items():
            self[k] = v

    def rename_key(self, old: str, new: str) -> bool:
        node = self._unwrap_root()
        T = JdaughterType
        if node.type != T.OBJECT:
            raise TypeError("rename_key() only for OBJECT")
        for e in node.children:
            if e.pair and e.pair.children[0].value == old:
                e.pair.children[0].value = new
                return True
        return False

    def move_key(self, key: str, new_index: int) -> bool:
        node = self._unwrap_root()
        T = JdaughterType
        if node.type != T.OBJECT:
            raise TypeError("move_key() only for OBJECT")
        entries = node.children
        for i, e in enumerate(entries):
            if e.pair and e.pair.children[0].value == key:
                entry = entries.pop(i)
                entries.insert(new_index, entry)
                return True
        return False

    # array helpers
    def append(self, value):
        node = self._unwrap_root()
        if node.type != JdaughterType.ARRAY:
            raise TypeError("append() only for ARRAY")
        node.children.append(JdaughterArrEntry(value=_wrap(value)))

    def insert(self, index: int, value):
        node = self._unwrap_root()
        if node.type != JdaughterType.ARRAY:
            raise TypeError("insert() only for ARRAY")
        node.children.insert(index, JdaughterArrEntry(value=_wrap(value)))

    def extend(self, iterable):
        node = self._unwrap_root()
        if node.type != JdaughterType.ARRAY:
            raise TypeError("extend() only for ARRAY")
        for v in iterable:
            node.children.append(JdaughterArrEntry(value=_wrap(v)))

    def pop_index(self, index: int = -1):
        node = self._unwrap_root()
        if node.type != JdaughterType.ARRAY:
            raise TypeError("pop_index() only for ARRAY")
        n = len(node.children)
        idx = index if index >= 0 else n + index
        if idx < 0 or idx >= n:
            raise IndexError(index)
        e = node.children.pop(idx)
        return e.value

    # path helper
    def ensure_path(self, *keys: str) -> "JdaughterNode":
        cur = self
        for k in keys:
            try:
                nxt = cur[k]
            except KeyError:
                cur[k] = {}
                nxt = cur[k]
            cur = nxt
        return cur

    # comment helpers
    def add_leading_comment(
        self, key: str, text: str, block: bool = False
    ) -> bool:
        node = self._unwrap_root()
        T = JdaughterType
        if node.type != T.OBJECT:
            raise TypeError("add_leading_comment() only for OBJECT")
        cnode = JdaughterNode(
            T.BLOCK_COMMENT if block else T.INLINE_COMMENT, text
        )
        for e in node.children:
            if e.pair and e.pair.children[0].value == key:
                e.leading.append(cnode)
                return True
        return False

    def add_trailing_comment(
        self, key: str, text: str, block: bool = False
    ) -> bool:
        node = self._unwrap_root()
        T = JdaughterType
        if node.type != T.OBJECT:
            raise TypeError("add_trailing_comment() only for OBJECT")
        cnode = JdaughterNode(
            T.BLOCK_COMMENT if block else T.INLINE_COMMENT, text
        )
        for e in node.children:
            if e.pair and e.pair.children[0].value == key:
                e.trailing.append(cnode)
                return True
        return False

    # python conversion
    def to_python(self):
        T = JdaughterType
        node = self._unwrap_root()
        if node.type == T.STRING:
            return node.value
        if node.type == T.NUMBER:
            return (
                float(node.value)
                if ("." in node.value or "e" in node.value or "E" in node.value)
                else int(node.value)
            )
        if node.type == T.BOOLEAN:
            return bool(node.value)
        if node.type == T.NULL:
            return None
        if node.type == T.NAN:
            return float("nan")
        if node.type == T.ARRAY:
            return [
                e.value.to_python()
                for e in node.children
                if e.value is not None
            ]
        if node.type == T.OBJECT:
            res = {}
            for e in node.children:
                if e.pair:
                    k = e.pair.children[0].value
                    v = e.pair.children[1].to_python()
                    res[k] = v
            return res
        if node.type == T.DOCUMENT:
            return node._unwrap_root().to_python()
        raise TypeError(f"to_python unsupported for {node.type}")

@dataclass
class JdaughterObjEntry:
    leading: List["JdaughterNode"] = field(default_factory=list)
    pair: Optional["JdaughterNode"] = None
    inline_after: List["JdaughterNode"] = field(default_factory=list)
    trailing: List["JdaughterNode"] = field(default_factory=list)

@dataclass
class JdaughterArrEntry:
    leading: List["JdaughterNode"] = field(default_factory=list)
    value: Optional["JdaughterNode"] = None
    inline_after: List["JdaughterNode"] = field(default_factory=list)
    trailing: List["JdaughterNode"] = field(default_factory=list)

# ---------- Parser ----------
class JdaughterParser:
    def __init__(self, toks: Iterable["JdaughterToken"]):
        self.toks = list(toks)
        self.i = 0
        self.last_tok: Optional["JdaughterToken"] = None

    def _p(self, k=0):
        return self.toks[self.i + k]

    def _adv(self):
        t = self._p()
        self.i += 1
        self.last_tok = t
        return t

    def _eat(self, tt):
        t = self._p()
        if t.type != tt:
            raise JdaughterParseError(f"expected {tt} got {t.type}")
        return self._adv()

    def _comment(self, t: "JdaughterToken") -> "JdaughterNode":
        return JdaughterNode(
            JdaughterType.INLINE_COMMENT
            if t.type == JdaughterTokType.INLINE_COMMENT
            else JdaughterType.BLOCK_COMMENT,
            t.lexeme,
        )

    def _collect_comments(self) -> List["JdaughterNode"]:
        out = []
        while self._p().type in (
            JdaughterTokType.INLINE_COMMENT,
            JdaughterTokType.BLOCK_COMMENT,
        ):
            out.append(self._comment(self._adv()))
        return out

    def parse(self) -> "JdaughterNode":
        pre = []
        while self._p().type in (
            JdaughterTokType.INLINE_COMMENT,
            JdaughterTokType.BLOCK_COMMENT,
        ):
            pre.append(self._comment(self._adv()))
        root = self._value()
        post = []
        while self._p().type in (
            JdaughterTokType.INLINE_COMMENT,
            JdaughterTokType.BLOCK_COMMENT,
        ):
            post.append(self._comment(self._adv()))
        self._eat(JdaughterTokType.EOF)
        return JdaughterNode(
            JdaughterType.DOCUMENT, None, [*pre, root, *post]
        )

    def _value(self) -> "JdaughterNode":
        t = self._p()
        TT = JdaughterTokType
        JT = JdaughterType
        if t.type == TT.LBRACE:
            return self._object()
        if t.type == TT.LBRACKET:
            return self._array()
        if t.type == TT.STRING:
            tok = self._adv()
            return JdaughterNode(JT.STRING, tok.lexeme)
        if t.type == TT.NUMBER:
            tok = self._adv()
            return JdaughterNode(JT.NUMBER, tok.lexeme)
        if t.type == TT.TRUE:
            self._adv()
            return JdaughterNode(JT.BOOLEAN, True)
        if t.type == TT.FALSE:
            self._adv()
            return JdaughterNode(JT.BOOLEAN, False)
        if t.type == TT.NULL:
            self._adv()
            return JdaughterNode(JT.NULL, None)
        if t.type == TT.NAN:
            self._adv()
            return JdaughterNode(JT.NAN, float("nan"))
        raise JdaughterParseError(f"unexpected {t.type}")

    def _object(self) -> "JdaughterNode":
        self._eat(JdaughterTokType.LBRACE)
        entries: List["JdaughterObjEntry"] = []

        pending_leading: List["JdaughterNode"] = self._collect_comments()

        if self._p().type == JdaughterTokType.RBRACE:
            self._adv()
            obj = JdaughterNode(JdaughterType.OBJECT, None, [])
            if pending_leading:
                obj.children.append(
                    JdaughterObjEntry(
                        leading=[],
                        pair=None,
                        trailing=pending_leading,
                    )
                )
            return obj

        while True:
            entry = JdaughterObjEntry()

            if pending_leading:
                entry.leading, pending_leading = pending_leading, []
            else:
                entry.leading = self._collect_comments()

            keytok = self._eat(JdaughterTokType.STRING)
            key = JdaughterNode(JdaughterType.KEY, keytok.lexeme)
            key_line = keytok.line

            between: List["JdaughterToken"] = []
            while self._p().type in (
                JdaughterTokType.INLINE_COMMENT,
                JdaughterTokType.BLOCK_COMMENT,
            ):
                between.append(self._adv())
            for tok in between:
                entry.leading.append(self._comment(tok))

            self._eat(JdaughterTokType.COLON)

            entry.leading.extend(self._collect_comments())

            val = self._value()
            val_line = self.last_tok.line if self.last_tok else key_line

            after_val: List["JdaughterToken"] = []
            while self._p().type in (
                JdaughterTokType.INLINE_COMMENT,
                JdaughterTokType.BLOCK_COMMENT,
            ):
                after_val.append(self._adv())

            near_val: List[tuple["JdaughterToken", "JdaughterNode"]] = []
            far_val_nodes: List["JdaughterNode"] = []
            for tok in after_val:
                node = self._comment(tok)
                if tok.line <= val_line + 1:
                    near_val.append((tok, node))
                else:
                    far_val_nodes.append(node)

            for tok, node in near_val:
                if (
                    tok.type == JdaughterTokType.INLINE_COMMENT
                    and tok.line == val_line
                ):
                    entry.inline_after.append(node)
                else:
                    entry.trailing.append(node)

            t = self._p()

            if t.type == JdaughterTokType.RBRACE:
                entry.trailing.extend(far_val_nodes)
                entry.pair = JdaughterNode(
                    JdaughterType.PAIR, None, [key, val]
                )
                entries.append(entry)
                self._adv()  # '}'
                break

            if t.type != JdaughterTokType.COMMA:
                raise JdaughterParseError(f"expected , or }} got {t.type}")

            comma_tok = self._adv()
            comma_line = comma_tok.line

            after_comma: List["JdaughterToken"] = []
            while self._p().type in (
                JdaughterTokType.INLINE_COMMENT,
                JdaughterTokType.BLOCK_COMMENT,
            ):
                after_comma.append(self._adv())

            near_comma: List[tuple["JdaughterToken", "JdaughterNode"]] = []
            far_comma_nodes: List["JdaughterNode"] = []
            for tok in after_comma:
                node = self._comment(tok)
                if tok.line <= comma_line + 1:
                    near_comma.append((tok, node))
                else:
                    far_comma_nodes.append(node)

            for tok, node in near_comma:
                if (
                    tok.type == JdaughterTokType.INLINE_COMMENT
                    and tok.line == comma_line
                ):
                    entry.inline_after.append(node)
                else:
                    entry.trailing.append(node)

            pending_leading = []
            pending_leading.extend(far_val_nodes)
            pending_leading.extend(far_comma_nodes)

            entry.pair = JdaughterNode(
                JdaughterType.PAIR, None, [key, val]
            )
            entries.append(entry)

            if (
                self._p().type == JdaughterTokType.RBRACE
                and not pending_leading
            ):
                self._adv()  # '}'
                break

        return JdaughterNode(JdaughterType.OBJECT, None, entries)

    def _array(self) -> "JdaughterNode":
        self._eat(JdaughterTokType.LBRACKET)
        entries: List["JdaughterArrEntry"] = []
        leading: List["JdaughterNode"] = self._collect_comments()

        if self._p().type == JdaughterTokType.RBRACKET:
            self._adv()
            arr = JdaughterNode(JdaughterType.ARRAY, None, [])
            if leading:
                arr.children.append(
                    JdaughterArrEntry(leading=[], value=None, trailing=leading)
                )
            return arr

        while True:
            e = JdaughterArrEntry()

            if leading:
                e.leading, leading = leading, []
            else:
                e.leading = self._collect_comments()

            val = self._value()
            e.value = val
            val_line = self.last_tok.line if self.last_tok else 0

            comments_after_val: List["JdaughterToken"] = []
            while self._p().type in (
                JdaughterTokType.INLINE_COMMENT,
                JdaughterTokType.BLOCK_COMMENT,
            ):
                comments_after_val.append(self._adv())

            t = self._p()

            def _attach_after_val_for_last():
                for tok in comments_after_val:
                    node = self._comment(tok)
                    if (
                        tok.type == JdaughterTokType.INLINE_COMMENT
                        and tok.line == val_line
                    ):
                        e.inline_after.append(node)
                    else:
                        e.trailing.append(node)

            if t.type == JdaughterTokType.COMMA:
                comma_tok = self._adv()

                same_line_inline_val: List["JdaughterNode"] = []
                other_val: List["JdaughterToken"] = []
                for tok in comments_after_val:
                    if (
                        tok.type == JdaughterTokType.INLINE_COMMENT
                        and tok.line == val_line
                    ):
                        same_line_inline_val.append(self._comment(tok))
                    else:
                        other_val.append(tok)

                e.inline_after.extend(same_line_inline_val)
                e.trailing.extend(self._comment(tok) for tok in other_val)

                comments_after_comma: List["JdaughterToken"] = []
                while self._p().type in (
                    JdaughterTokType.INLINE_COMMENT,
                    JdaughterTokType.BLOCK_COMMENT,
                ):
                    comments_after_comma.append(self._adv())

                same_line_inline_comma: List["JdaughterNode"] = []
                after_tokens: List["JdaughterToken"] = []
                for tok in comments_after_comma:
                    if (
                        tok.type == JdaughterTokType.INLINE_COMMENT
                        and tok.line == comma_tok.line
                    ):
                        same_line_inline_comma.append(self._comment(tok))
                    else:
                        after_tokens.append(tok)

                e.inline_after.extend(same_line_inline_comma)

                trailing_tokens: List["JdaughterToken"] = []
                leading_tokens: List["JdaughterToken"] = []
                if after_tokens:
                    trailing_tokens.append(after_tokens[0])
                    last_line = after_tokens[0].line
                    idx = 1
                    while idx < len(after_tokens):
                        tok = after_tokens[idx]
                        if tok.line == last_line + 1:
                            trailing_tokens.append(tok)
                            last_line = tok.line
                            idx += 1
                        else:
                            break
                    leading_tokens = after_tokens[idx:]

                e.trailing.extend(self._comment(tok) for tok in trailing_tokens)
                leading = [self._comment(tok) for tok in leading_tokens]

                if self._p().type == JdaughterTokType.RBRACKET and leading:
                    e.trailing.extend(leading)
                    leading = []

                entries.append(e)

                if self._p().type == JdaughterTokType.RBRACKET:
                    self._adv()
                    break

            elif t.type == JdaughterTokType.RBRACKET:
                
                _attach_after_val_for_last()
                entries.append(e)
                self._adv()
                break
            else:
                raise JdaughterParseError(f"expected , or ] got {t.type}")

        return JdaughterNode(JdaughterType.ARRAY, None, entries)

# ---------- Dumper ----------
def _esc(s: str) -> str:
    rep = {
        "\\": "\\\\",
        '"': '\\"',
        "\b": "\\b",
        "\f": "\\f",
        "\n": "\\n",
        "\r": "\\r",
        "\t": "\\t",
    }
    out = []
    for ch in s:
        if ch in rep:
            out.append(rep[ch])
        elif ord(ch) < 0x20:
            out.append(f"\\u{ord(ch):04x}")
        else:
            out.append(ch)
    return '"' + "".join(out) + '"'

def dumps(node: "JdaughterNode", indent: int = 2) -> str:
    out: List[str] = []

    def w(s: str):
        out.append(s)

    def nl(level: int):
        if indent is not None:
            w("\n" + " " * (indent * level))

    def emit_comment(n: "JdaughterNode"):
        if n.type == JdaughterType.INLINE_COMMENT:
            w("//" + n.value)
        else:
            w("/*" + n.value + "*/")

    def emit(n: "JdaughterNode", level: int):
        t = n.type
        if t == JdaughterType.DOCUMENT:
            for ch in n.children:
                if isinstance(ch, JdaughterNode) and ch.type in (
                    JdaughterType.INLINE_COMMENT,
                    JdaughterType.BLOCK_COMMENT,
                ):
                    emit_comment(ch)
                    nl(0)
                else:
                    emit(ch, level)
            return
        if t == JdaughterType.STRING:
            w(_esc(n.value))
            return
        if t == JdaughterType.NUMBER:
            w(n.value)
            return
        if t == JdaughterType.BOOLEAN:
            w("true" if n.value else "false")
            return
        if t == JdaughterType.NULL:
            w("null")
            return
        if t == JdaughterType.NAN:
            w("NaN")
            return
        if t == JdaughterType.KEY:
            w(_esc(n.value))
            return
        if t == JdaughterType.PAIR:
            k, v = n.children
            emit(k, level)
            w(": ")
            emit(v, level)
            return

        if t == JdaughterType.OBJECT:
            w("{")
            level += 1
            nl(level)
            entries: List["JdaughterObjEntry"] = n.children
            for idx, e in enumerate(entries):
                assert isinstance(
                    e, JdaughterObjEntry
                ), f"OBJECT child must be JdaughterObjEntry, got {type(e).__name__}"
                for c in e.leading:
                    emit_comment(c)
                    nl(level)

                if e.pair is not None:
                    emit(e.pair, level)
                    
                    for c in e.inline_after:
                        w(" ")
                        emit_comment(c)
                    
                    for c in e.trailing:
                        nl(level)
                        emit_comment(c)

                if idx != len(entries) - 1:
                    w(",")
                    nl(level)

            level -= 1
            if entries:
                nl(level)
            w("}")
            return

        if t == JdaughterType.ARRAY:
            w("[")
            level += 1
            nl(level)
            entries: List["JdaughterArrEntry"] = n.children
            for idx, e in enumerate(entries):
                assert isinstance(
                    e, JdaughterArrEntry
                ), f"ARRAY child must be JdaughterArrEntry, got {type(e).__name__}"
                for c in e.leading:
                    emit_comment(c)
                    nl(level)

                if e.value is not None:
                    emit(e.value, level)
                    for c in e.inline_after:
                        w(" ")
                        emit_comment(c)
                    for c in e.trailing:
                        nl(level)
                        emit_comment(c)

                if idx != len(entries) - 1:
                    w(",")
                    nl(level)
            level -= 1
            if entries:
                nl(level)
            w("]")
            return

        raise TypeError(f"unknown {t}")

    emit(node, 0)
    return "".join(out)

def _wrap(value) -> "JdaughterNode":
    if isinstance(value, JdaughterNode):
        return value
    JT = JdaughterType
    if value is None:
        return JdaughterNode(JT.NULL, None)
    if isinstance(value, bool):
        return JdaughterNode(JT.BOOLEAN, bool(value))
    if isinstance(value, str):
        return JdaughterNode(JT.STRING, value)
    if isinstance(value, (int, float)):
        return JdaughterNode(JT.NUMBER, repr(value))
    if isinstance(value, list):
        entries = [JdaughterArrEntry(value=_wrap(v)) for v in value]
        return JdaughterNode(JT.ARRAY, None, entries)
    if isinstance(value, dict):
        entries = [
            JdaughterObjEntry(pair=make_pair(k, _wrap(v)))
            for k, v in value.items()
        ]
        return JdaughterNode(JT.OBJECT, None, entries)
    raise TypeError(f"cannot wrap to JdaughterNode: {type(value).__name__}")

# ---------- Public API ----------
def loads(s: str) -> "JdaughterNode":
    lexer = JdaughterLexer(s)
    parser = JdaughterParser(lexer.tokens())
    return parser.parse()

def dump(node: "JdaughterNode", fp, indent: int = 2) -> None:
    fp.write(dumps(node, indent))

def load(fp) -> "JdaughterNode":
    return loads(fp.read())

# ---------- Helpers ----------
def comment_line(text: str) -> "JdaughterNode":
    return JdaughterNode(JdaughterType.INLINE_COMMENT, text)

def comment_block(text: str) -> "JdaughterNode":
    return JdaughterNode(JdaughterType.BLOCK_COMMENT, text)

def make_pair(key: str, value: "JdaughterNode") -> "JdaughterNode":
    return JdaughterNode(
        JdaughterType.PAIR,
        None,
        [JdaughterNode(JdaughterType.KEY, key), value],
    )

__all__ = [
    "JdaughterTokType",
    "JdaughterToken",
    "JdaughterLexerError",
    "JdaughterParseError",
    "JdaughterType",
    "JdaughterNode",
    "JdaughterObjEntry",
    "JdaughterArrEntry",
    "loads",
    "dumps",
    "load",
    "dump",
    "comment_line",
    "comment_block",
    "make_pair",
]