#!/usr/bin/env python3
"""
JsonSAT: Stateful JSON constraint for Helix64 token space.

Provides streaming JSON parser that works on token IDs (0-63) to enforce
valid JSON structure during generation. ~100 LOC as requested.

Token mapping:
  0-5:   Structural ({, }, [, ], :, ,)
  6-9:   Literals (null, true, false, ")
  10-15: Whitespace + escape sequences
  16-25: Digits (0-9)
  26-51: Letters (a-z uppercase/lowercase)
  52-63: Special chars for numbers/strings

States:
  - START: Expect { or [
  - OBJECT_START: Expect " (key) or }
  - OBJECT_KEY: Building key string
  - OBJECT_COLON: Expect :
  - OBJECT_VALUE: Expect value (string, number, object, array, literal)
  - OBJECT_COMMA: Expect , or }
  - ARRAY_START: Expect value or ]
  - ARRAY_VALUE: Building value
  - ARRAY_COMMA: Expect , or ]
  - IN_STRING: Building string content
  - IN_NUMBER: Building number
"""

import hashlib

# JSON vocabulary (64 tokens for Helix64 compatibility)
JSON_VOCAB = [
    # Structural (0-5)
    "{",      # 0
    "}",      # 1
    "[",      # 2
    "]",      # 3
    ":",      # 4
    ",",      # 5

    # Literals (6-9)
    "null",   # 6
    "true",   # 7
    "false",  # 8
    '"',      # 9 (quote for strings)

    # Whitespace + escapes (10-15)
    " ",      # 10
    "\n",     # 11
    "\t",     # 12
    "\\",     # 13 (backslash for escapes)
    "/",      # 14 (forward slash)
    ".",      # 15 (decimal point)

    # Digits (16-25)
    "0", "1", "2", "3", "4", "5", "6", "7", "8", "9",

    # Lowercase letters (26-51)
    "a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m",
    "n", "o", "p", "q", "r", "s", "t", "u", "v", "w", "x", "y", "z",

    # Uppercase letters + punctuation (52-63) - allows mixed case and common punctuation
    "A", "B", "C", "D", "E", "F", "G", "H", "I", "J", "_", "-",
]

# Pinned vocab hash (known-good v2 with _/- at pos 62-63)
JSON_VOCAB_SHA256_EXPECTED = "19b339c7a0f028374d75939f8a91d73ee51bea1c071dce5c994bc74f3ef1554c"
JSON_VOCAB_VERSION = "json_vocab.v2"  # Schema version for receipts (v2 = with _/- at pos 62-63)

# Structural invariants (catch accidental edits at import time)
assert len(JSON_VOCAB) == 64, f"JSON_VOCAB length must be 64, got {len(JSON_VOCAB)}"
assert JSON_VOCAB[62] == "_", f"JSON_VOCAB[62] must be '_', got {JSON_VOCAB[62]!r}"
assert JSON_VOCAB[63] == "-", f"JSON_VOCAB[63] must be '-', got {JSON_VOCAB[63]!r}"
assert len(set(JSON_VOCAB)) == 64, f"JSON_VOCAB has duplicates: {len(set(JSON_VOCAB))} unique of 64"

# Hash check: compare computed vs pinned (detects accidental vocab edits)
_COMPUTED_HASH = hashlib.sha256("".join(JSON_VOCAB).encode()).hexdigest()
assert _COMPUTED_HASH == JSON_VOCAB_SHA256_EXPECTED, (
    f"JSON_VOCAB hash mismatch: computed={_COMPUTED_HASH[:16]}... expected={JSON_VOCAB_SHA256_EXPECTED[:16]}..."
)

# Export alias for cdc_decode compatibility
JSON_VOCAB_SHA256 = JSON_VOCAB_SHA256_EXPECTED

def json_id_of(tok: str) -> int:
    """Get token ID for a JSON token (name-based lookup)."""
    try:
        return JSON_VOCAB.index(tok)
    except ValueError:
        raise ValueError(f"Token {tok!r} not in JSON_VOCAB")

def json_id_to_token(i: int) -> str:
    """Convert token ID to JSON token string."""
    return JSON_VOCAB[i % 64]

def ids_to_json(ids) -> str:
    """Convert token IDs to JSON string."""
    return "".join(JSON_VOCAB[i % 64] for i in ids)


def json_token_class(i: int) -> str:
    """
    Get token class for JSON token ID (for scoring stability).

    Classes:
    - STRUCT: {, }, [, ], :, ,
    - LITERAL: null, true, false
    - STRING: "
    - WS: space, newline, tab
    - DIGIT: 0-9
    - LETTER: a-z, A-L
    - OTHER: \, /, .
    """
    tok = JSON_VOCAB[i % 64]

    if tok in ("{", "}", "[", "]", ":", ","):
        return "STRUCT"
    elif tok in ("null", "true", "false"):
        return "LITERAL"
    elif tok == '"':
        return "STRING"
    elif tok in (" ", "\n", "\t"):
        return "WS"
    elif tok.isdigit():
        return "DIGIT"
    elif tok.isalpha() and len(tok) == 1:
        return "LETTER"
    else:
        return "OTHER"


class JsonSAT:
    """
    Stateful JSON constraint engine for Helix64 tokens.

    Tracks JSON structure (object/array stack) and returns allowed token IDs
    based on current parse state. Enforces valid JSON syntax at token level.
    """

    def __init__(self):
        self.state = "START"
        self.stack = []  # Track object/array nesting
        self.in_string = False
        self.string_escaped = False
        self.hex_count = 0  # For \uXXXX sequences
        self.hex_buffer = []  # Buffer to collect hex digits for surrogate check
        self.depth = 0  # Track nesting depth for determinism salt
        self.num_state = "NONE"  # NONE|INT|FRAC|EXP|EXP_SIGN
        self.ws_run = 0  # consecutive raw spaces (outside strings)
        self.sur_hi_pending = False  # High surrogate must be followed by low surrogate
        self.num_leading_zero = False  # Track leading zero to block 012

    def consume(self, tok: str):
        """Advance internal state as if `tok` was just emitted."""
        # Re-run a single-step allow to apply transitions deterministically.
        # This keeps state in sync so the decoder can query terminal immediately.
        _ = self.allowed_ids([tok])

    def closing_ids(self) -> list:
        """Return closers valid in current context (may be empty)."""
        closers = []
        if self.stack:
            top = self.stack[-1]
            if top == "object":
                closers.append(json_id_of("}"))
            elif top == "array":
                closers.append(json_id_of("]"))
        return closers

    def min_close_steps(self) -> int:
        """Crude lower bound on tokens to finish from here (for endgame clamp)."""
        steps = self.depth
        if self.in_string:
            steps += 1  # need closing quote
        # If expecting a value at top of an object/array, assume 'null' (1 token)
        if self.state in ("OBJECT_START", "ARRAY_START", "OBJECT_COLON"):
            steps += 1
        return max(0, steps)

    def allowed_ids(self, tail_tokens: list) -> list:
        """
        Get allowed token IDs based on current JSON state.

        Args:
            tail_tokens: Last N generated tokens (for context)

        Returns:
            List of allowed token IDs for next position
        """
        if not tail_tokens:
            # Start: only allow { or [
            return [json_id_of("{"), json_id_of("[")]

        last = tail_tokens[-1] if tail_tokens else ""

        # String mode: handle escape sequences
        if self.in_string:
            # High surrogate pending: next token MUST be \uXXXX with low surrogate
            if self.sur_hi_pending and not self.string_escaped:
                if last == "\\":
                    self.string_escaped = True
                    # Only allow 'u' after backslash
                    return [json_id_of("u")]
                else:
                    # Invalid: high surrogate not followed by escape
                    # POLICY: Terminate string immediately (strict mode)
                    # Keep sur_hi_pending=True to force quote (don't silently clear)
                    return [json_id_of('"')]

            # \uXXXX hex escape sequence (already in hex mode)
            if self.hex_count > 0:
                # Collect hex digit into buffer
                if self.hex_count <= 4 and last in "0123456789abcdefABCDEF":
                    self.hex_buffer.append(last)

                self.hex_count += 1
                if self.hex_count > 4:
                    # Completed 4 hex digits, check for surrogate pairs
                    self.hex_count = 0
                    hex_str = "".join(self.hex_buffer[-4:]) if len(self.hex_buffer) >= 4 else "0000"
                    try:
                        codepoint = int(hex_str, 16)
                        # High surrogate range: D800-DBFF
                        if 0xD800 <= codepoint <= 0xDBFF:
                            self.sur_hi_pending = True
                            self.hex_buffer = []  # Clear buffer after use
                            # Must be followed by \uDC00-DFFF
                            return [json_id_of("\\")]
                        # Low surrogate range: DC00-DFFF
                        elif 0xDC00 <= codepoint <= 0xDFFF:
                            if self.sur_hi_pending:
                                # Valid: high followed by low
                                self.sur_hi_pending = False
                            else:
                                # Invalid: lone low surrogate (but allow for now)
                                pass
                        else:
                            # Normal codepoint, clear any pending
                            self.sur_hi_pending = False
                    except Exception:
                        self.sur_hi_pending = False

                    # Clear hex buffer after processing codepoint
                    self.hex_buffer = []

                    # Return to normal string mode
                    allowed = [json_id_of(" "), json_id_of('"'), json_id_of("\\")]
                    allowed += [json_id_of(c) for c in "abcdefghijklmnopqrstuvwxyz"]
                    allowed += [json_id_of(c) for c in "ABCDEFGHIJ"]
                    allowed += [json_id_of(str(d)) for d in range(10)]
                    SAFE_PUNCT = "/._-:,"
                    allowed += [json_id_of(c) for c in SAFE_PUNCT]
                    return allowed
                else:
                    # Need more hex digits (0-9, a-f, A-F)
                    allowed = [json_id_of(str(d)) for d in range(10)]
                    allowed += [json_id_of(c) for c in "abcdef"]
                    allowed += [json_id_of(c) for c in "ABCDEF"]
                    return allowed

            # After \u: start hex sequence (check BEFORE clearing string_escaped)
            if last == "u" and self.string_escaped:
                self.string_escaped = False
                self.hex_count = 1
                self.hex_buffer = []  # Clear buffer when starting new \uXXXX
                # Need hex digits
                allowed = [json_id_of(str(d)) for d in range(10)]
                allowed += [json_id_of(c) for c in "abcdef"]
                allowed += [json_id_of(c) for c in "ABCDEF"]
                return allowed

            # After backslash: allow escape sequences
            if self.string_escaped:
                self.string_escaped = False
                # Valid JSON escapes: " \ / b f n r t u
                allowed = [json_id_of('"'), json_id_of("\\"), json_id_of("/")]
                allowed += [json_id_of(c) for c in "bfnrtu"]
                return allowed

            # Backslash: enter escape mode
            elif last == "\\":
                self.string_escaped = True
                return [json_id_of('"'), json_id_of("\\"), json_id_of("/"),
                        json_id_of("b"), json_id_of("f"), json_id_of("n"),
                        json_id_of("r"), json_id_of("t"), json_id_of("u")]

            # End string
            elif last == '"':
                self.in_string = False
                self.hex_buffer = []  # Clear buffer when exiting string mode
                self.sur_hi_pending = False  # Clear pending flag when string closes
                # Determine what's allowed after string
                if self.state == "OBJECT_KEY":
                    self.state = "OBJECT_COLON"
                    return [json_id_of(":"), json_id_of(" ")]
                elif self.state == "OBJECT_VALUE":
                    self.state = "OBJECT_COMMA"
                    return [json_id_of(","), json_id_of("}"), json_id_of(" ")]
                elif self.state == "ARRAY_VALUE":
                    self.state = "ARRAY_COMMA"
                    return [json_id_of(","), json_id_of("]"), json_id_of(" ")]
            else:
                # Inside string: allow letters, digits, space, backslash, safe punctuation
                allowed = [json_id_of(" "), json_id_of('"'), json_id_of("\\")]
                allowed += [json_id_of(c) for c in "abcdefghijklmnopqrstuvwxyz"]
                allowed += [json_id_of(c) for c in "ABCDEFGHIJ"]
                allowed += [json_id_of(str(d)) for d in range(10)]
                # Safe punctuation (newline must be escaped, but these are fine)
                SAFE_PUNCT = "/._-:,"
                allowed += [json_id_of(c) for c in SAFE_PUNCT]
                return allowed

        # Outside string: track raw-space runs to avoid drifting on whitespace
        if last == " ":
            self.ws_run = min(2, self.ws_run + 1)
        else:
            self.ws_run = 0

        # JSON-number enforcement (tighter than RFC 8259):
        # - No leading zeros except "0" alone (blocks "012")
        # - No "+" sign (not in 64-token vocab)
        # - Lowercase 'e' only for exponent (uppercase 'E' omitted by design; vocab excludes it)
        # - No trailing decimal point (e.g., "1." rejected)
        # - No leading decimal point (e.g., ".5" rejected)
        # Minimal number sub-FSM (prevents endless digits / half-numbers)
        def _num_allowed():
            if self.num_state == "NONE":
                # entering number only if last was digit or '-' (already emitted)
                allowed = [json_id_of(str(d)) for d in range(10)] + [json_id_of("."),
                       json_id_of(","), json_id_of("]"), json_id_of("}"), json_id_of(" ")]
                # Leading zero rule: if last was '0', block additional digits
                if self.num_leading_zero:
                    allowed = [i for i in allowed if json_id_to_token(i) not in "123456789"]
                return allowed
            if self.num_state == "INT":
                allowed = [json_id_of(str(d)) for d in range(10)] + [json_id_of("."),
                       json_id_of(","), json_id_of("]"), json_id_of("}"), json_id_of(" ")]
                # Allow lowercase 'e' for exponent (uppercase 'E' omitted; not in vocab)
                try:
                    allowed += [json_id_of("e")]
                except Exception:
                    pass
                # Leading zero rule: if num_leading_zero, block digits (only allow '.', 'e', structural)
                if self.num_leading_zero:
                    allowed = [i for i in allowed if json_id_to_token(i) not in "123456789"]
                return allowed
            if self.num_state == "FRAC":
                allowed = [json_id_of(str(d)) for d in range(10)] + [
                       json_id_of(","), json_id_of("]"), json_id_of("}"), json_id_of(" ")]
                # Allow lowercase 'e' for exponent (uppercase 'E' omitted; not in vocab)
                try:
                    allowed += [json_id_of("e")]
                except Exception:
                    pass
                return allowed
            if self.num_state == "EXP":
                # '+' not in JSON_VOCAB, so only allow '-' or digits
                return [json_id_of("-")] + [json_id_of(str(d)) for d in range(10)]
            if self.num_state == "EXP_SIGN":
                return [json_id_of(str(d)) for d in range(10)] + [
                       json_id_of(","), json_id_of("]"), json_id_of("}"), json_id_of(" ")]
            return []

        # Update num_state by inspecting `last`
        if last.isdigit():
            # Track leading zero: if first digit is '0', set flag
            if self.num_state == "NONE" and last == "0":
                self.num_leading_zero = True
            elif last != "0":
                self.num_leading_zero = False

            if self.num_state in ("NONE", "INT"):
                self.num_state = "INT"
            elif self.num_state in ("FRAC", "EXP", "EXP_SIGN"):
                # If EXP or EXP_SIGN, digits keep us there until a non-digit lands
                pass
        elif last == "." and self.num_state in ("NONE", "INT"):
            self.num_state = "FRAC"
            self.num_leading_zero = False  # '0.' is valid, clear flag
        elif last in ("e",) and self.num_state in ("INT", "FRAC"):
            self.num_state = "EXP"
            self.num_leading_zero = False
        elif last in ("+", "-") and self.num_state == "EXP":
            self.num_state = "EXP_SIGN"
        elif last == "-" and self.num_state == "NONE":
            # Negative number start, keep num_state = NONE until digit appears
            pass
        else:
            # Any structural/non-number token resets number FSM
            self.num_state = "NONE"
            self.num_leading_zero = False

        # State: START
        if self.state == "START":
            if last == "{":
                self.state = "OBJECT_START"
                self.stack.append("object")
                self.depth += 1
                return [json_id_of('"'), json_id_of("}"), json_id_of(" ")]
            elif last == "[":
                self.state = "ARRAY_START"
                self.stack.append("array")
                self.depth += 1
                return self._value_tokens() + [json_id_of("]"), json_id_of(" ")]

        # State: OBJECT_START
        if self.state == "OBJECT_START":
            if last == '"':
                self.state = "OBJECT_KEY"
                self.in_string = True
                # Start building key
                return [json_id_of(c) for c in "abcdefghijklmnopqrstuvwxyz"]
            elif last == "}":
                # Empty object
                return self._after_value()
            else:
                # Still waiting for key or close - allow space drift but not wrong brackets
                return [json_id_of('"'), json_id_of("}"), json_id_of(" ")]

        # State: OBJECT_COLON
        if self.state == "OBJECT_COLON":
            if last == ":":
                self.state = "OBJECT_VALUE"
                return self._value_tokens() + [json_id_of(" ")]

        # State: OBJECT_VALUE (start of value)
        if self.state == "OBJECT_VALUE":
            if last == '"':
                self.in_string = True
                return [json_id_of(c) for c in "abcdefghijklmnopqrstuvwxyz"]
            elif last in ("null", "true", "false"):
                self.state = "OBJECT_COMMA"
                return [json_id_of(","), json_id_of("}"), json_id_of(" ")]
            elif last == "{":
                self.state = "OBJECT_START"
                self.stack.append("object")
                self.depth += 1
                return [json_id_of('"'), json_id_of("}"), json_id_of(" ")]
            elif last == "[":
                self.state = "ARRAY_START"
                self.stack.append("array")
                self.depth += 1
                return self._value_tokens() + [json_id_of("]"), json_id_of(" ")]
            elif last == "-":
                # After minus sign: only allow digits (no '.' to prevent '-.5')
                return [json_id_of(str(d)) for d in range(10)]
            elif last.isdigit() or last == ".":
                # Building number - stay in OBJECT_VALUE, consult number FSM
                allowed = [json_id_of(str(d)) for d in range(10)] + [json_id_of("."), json_id_of(","), json_id_of("}"), json_id_of(" ")]
                # Allow 'e' for exponent if in INT or FRAC state
                if self.num_state in ("INT", "FRAC"):
                    try:
                        allowed.append(json_id_of("e"))
                    except Exception:
                        pass
                # Apply leading-zero rule if flag is set
                if self.num_leading_zero:
                    allowed = [i for i in allowed if json_id_to_token(i) not in "123456789"]
                return allowed
            elif last == "e":
                # In exponent: allow '-' or digits
                if self.num_state == "EXP":
                    return [json_id_of("-")] + [json_id_of(str(d)) for d in range(10)]
                # Fall through if not in EXP state
            elif last == ",":
                # Just finished a number value with comma - transition to OBJECT_AFTER_COMMA
                self.state = "OBJECT_AFTER_COMMA"
                return [json_id_of('"'), json_id_of(" ")]
            elif last == "}":
                # Just finished a number value with close - pop and continue
                return self._after_value()

        # State: OBJECT_COMMA
        if self.state == "OBJECT_COMMA":
            if last == ",":
                # after comma we require a new key string, not a closer
                self.state = "OBJECT_AFTER_COMMA"
                return [json_id_of('"'), json_id_of(" ")]
            elif last == "}":
                return self._after_value()
            else:
                return [json_id_of(","), json_id_of("}"), json_id_of(" ")]

        # State: OBJECT_AFTER_COMMA
        if self.state == "OBJECT_AFTER_COMMA":
            # only a key string may begin here (spaces are okay)
            if last == '"':
                self.state = "OBJECT_KEY"
                self.in_string = True
                # Return in-string allowed tokens
                allowed = [json_id_of(" "), json_id_of('"'), json_id_of("\\")]
                allowed += [json_id_of(c) for c in "abcdefghijklmnopqrstuvwxyz"]
                allowed += [json_id_of(c) for c in "ABCDEFGHIJ"]
                allowed += [json_id_of(str(d)) for d in range(10)]
                SAFE_PUNCT = "/._-:,"
                allowed += [json_id_of(c) for c in SAFE_PUNCT]
                return allowed
            # allow spaces to drift, but not '}'
            return [json_id_of('"'), json_id_of(" ")]

        # State: ARRAY_START
        if self.state == "ARRAY_START":
            if last == '"':
                self.state = "ARRAY_VALUE"
                self.in_string = True
                return [json_id_of(c) for c in "abcdefghijklmnopqrstuvwxyz"]
            elif last in ("null", "true", "false"):
                self.state = "ARRAY_COMMA"
                return [json_id_of(","), json_id_of("]"), json_id_of(" ")]
            elif last == "{":
                self.state = "OBJECT_START"
                self.stack.append("object")
                self.depth += 1
                return [json_id_of('"'), json_id_of("}"), json_id_of(" ")]
            elif last == "[":
                self.stack.append("array")
                self.depth += 1
                return self._value_tokens() + [json_id_of("]"), json_id_of(" ")]
            elif last == "-":
                # After minus sign: only allow digits (no '.' to prevent '-.5')
                self.state = "ARRAY_VALUE"
                return [json_id_of(str(d)) for d in range(10)]
            elif last.isdigit() or last == ".":
                # Building number in array
                self.state = "ARRAY_VALUE"
                allowed = [json_id_of(str(d)) for d in range(10)] + [json_id_of("."), json_id_of(","), json_id_of("]"), json_id_of(" ")]
                # Allow 'e' for exponent if in INT or FRAC state
                if self.num_state in ("INT", "FRAC"):
                    try:
                        allowed.append(json_id_of("e"))
                    except Exception:
                        pass
                # Apply leading-zero rule if flag is set
                if self.num_leading_zero:
                    allowed = [i for i in allowed if json_id_to_token(i) not in "123456789"]
                return allowed
            elif last == "e":
                # In exponent: allow '-' or digits
                if self.num_state == "EXP":
                    self.state = "ARRAY_VALUE"
                    return [json_id_of("-")] + [json_id_of(str(d)) for d in range(10)]
                # Fall through if not in EXP state
            elif last == "]":
                # Empty array
                return self._after_value()
            else:
                # Initial state - allow any value or close
                return self._value_tokens() + [json_id_of("]"), json_id_of(" ")]

        # State: ARRAY_VALUE
        if self.state == "ARRAY_VALUE":
            if last == '"':
                # Already handled by in_string logic above
                pass
            elif last == "-":
                # After minus sign: only allow digits (no '.' to prevent '-.5')
                return [json_id_of(str(d)) for d in range(10)]
            elif last.isdigit() or last == ".":
                # Building number - stay in ARRAY_VALUE, consult number FSM
                allowed = [json_id_of(str(d)) for d in range(10)] + [json_id_of("."), json_id_of(","), json_id_of("]"), json_id_of(" ")]
                # Allow 'e' for exponent if in INT or FRAC state
                if self.num_state in ("INT", "FRAC"):
                    try:
                        allowed.append(json_id_of("e"))
                    except Exception:
                        pass
                # Apply leading-zero rule if flag is set
                if self.num_leading_zero:
                    allowed = [i for i in allowed if json_id_to_token(i) not in "123456789"]
                return allowed
            elif last == "e":
                # In exponent: allow '-' or digits
                if self.num_state == "EXP":
                    return [json_id_of("-")] + [json_id_of(str(d)) for d in range(10)]
                # Fall through if not in EXP state
            elif last == ",":
                # Just finished a value with comma - transition to ARRAY_AFTER_COMMA
                self.state = "ARRAY_AFTER_COMMA"
                return self._value_tokens() + [json_id_of(" ")]
            elif last == "]":
                # Just finished a value with close - pop and continue
                return self._after_value()

        # State: ARRAY_COMMA
        if self.state == "ARRAY_COMMA":
            if last == ",":
                # must start another value; cannot close immediately
                self.state = "ARRAY_AFTER_COMMA"
                return self._value_tokens() + [json_id_of(" ")]
            elif last == "]":
                return self._after_value()
            else:
                return [json_id_of(","), json_id_of("]"), json_id_of(" ")]

        # State: ARRAY_AFTER_COMMA
        if self.state == "ARRAY_AFTER_COMMA":
            # require a value; no immediate closing bracket
            if last == '"':
                self.state = "ARRAY_VALUE"
                self.in_string = True
                return [json_id_of(c) for c in "abcdefghijklmnopqrstuvwxyz"]
            elif last in ("null", "true", "false"):
                self.state = "ARRAY_COMMA"
                return [json_id_of(","), json_id_of("]"), json_id_of(" ")]
            elif last == "{":
                self.state = "OBJECT_START"
                self.stack.append("object")
                self.depth += 1
                return [json_id_of('"'), json_id_of("}"), json_id_of(" ")]
            elif last == "[":
                self.stack.append("array")
                self.depth += 1
                return self._value_tokens() + [json_id_of("]"), json_id_of(" ")]
            elif last == "-":
                # After minus sign: only allow digits (no '.' to prevent '-.5')
                self.state = "ARRAY_VALUE"
                return [json_id_of(str(d)) for d in range(10)]
            elif last.isdigit() or last == ".":
                # Building number after comma
                self.state = "ARRAY_VALUE"
                allowed = [json_id_of(str(d)) for d in range(10)] + [json_id_of("."), json_id_of(","), json_id_of("]"), json_id_of(" ")]
                # Allow 'e' for exponent if in INT or FRAC state
                if self.num_state in ("INT", "FRAC"):
                    try:
                        allowed.append(json_id_of("e"))
                    except Exception:
                        pass
                # Apply leading-zero rule if flag is set
                if self.num_leading_zero:
                    allowed = [i for i in allowed if json_id_to_token(i) not in "123456789"]
                return allowed
            elif last == "e":
                # In exponent: allow '-' or digits
                if self.num_state == "EXP":
                    self.state = "ARRAY_VALUE"
                    return [json_id_of("-")] + [json_id_of(str(d)) for d in range(10)]
                # Fall through if not in EXP state
            else:
                # Allow any value start
                return self._value_tokens() + [json_id_of(" ")]

        # Fallback: state-aware minimal structural tokens (never introduce commas)
        allowed = [json_id_of("}"), json_id_of("]"), json_id_of(" ")]

        # Outside-string whitespace throttling (don't let space repeat forever)
        if self.ws_run >= 1 and json_id_of(" ") in allowed:
            allowed = [i for i in allowed if i != json_id_of(" ")]

        # If we are in a numeric segment, intersect with numeric allowances
        if self.num_state != "NONE":
            num_allowed = set(_num_allowed())
            allowed = [i for i in allowed if i in num_allowed]

        return allowed

    def _value_tokens(self) -> list:
        """Get token IDs for value start (string, number, literal, object, array)."""
        vals = [
            json_id_of('"'),      # string
            json_id_of("{"),      # object
            json_id_of("["),      # array
            json_id_of("null"),   # null
            json_id_of("true"),   # true
            json_id_of("false"),  # false
        ] + [json_id_of(str(d)) for d in range(10)]  # numbers
        # permit negative numbers at value start
        try:
            vals.append(json_id_of("-"))
        except Exception:
            pass
        return vals

    def _after_value(self) -> list:
        """Get allowed tokens after completing a value."""
        if not self.stack:
            # Already at top level (shouldn't happen, but guard)
            return []

        container = self.stack.pop()
        self.depth -= 1  # Decrement depth when closing container

        # Check if we're now at top level (closed the root container)
        if not self.stack:
            # Top level complete - set terminal state
            if container == "object":
                self.state = "OBJECT_COMMA"
            elif container == "array":
                self.state = "ARRAY_COMMA"
            return []  # No more tokens allowed

        # Still nested - continue based on parent container
        if container == "object":
            self.state = "OBJECT_COMMA"
            return [json_id_of(","), json_id_of("}"), json_id_of(" ")]
        elif container == "array":
            self.state = "ARRAY_COMMA"
            return [json_id_of(","), json_id_of("]"), json_id_of(" ")]
        return []

    def is_terminal(self) -> bool:
        """Check if parser is in a terminal state (valid end position)."""
        # START is NOT terminal - we need actual output
        return len(self.stack) == 0 and self.state in ("OBJECT_COMMA", "ARRAY_COMMA")

    def widen(self, tail_tokens: list) -> list:
        """
        Widen allowed set to prevent wedging.

        Returns a small, valid superset when allowed_ids would be empty.
        This is a safety fallback - should rarely trigger if FSM is correct.
        """
        # If we're in terminal state, don't widen - we're done
        if self.is_terminal():
            return []

        # If in string, allow basic string content
        if self.in_string:
            allowed = [json_id_of(" "), json_id_of('"')]
            allowed += [json_id_of(c) for c in "abcdefghijklmnopqrstuvwxyz"]
            allowed += [json_id_of(str(d)) for d in range(10)]
            return allowed

        # If in container, allow structural tokens + values based on state
        if self.stack:
            container = self.stack[-1]
            if container == "object":
                if self.state in ("OBJECT_START",):
                    # first pair or close — no comma allowed
                    return [json_id_of('"'), json_id_of("}"), json_id_of(" ")]
                if self.state in ("OBJECT_AFTER_COMMA",):
                    # must begin a key — no close here
                    return [json_id_of('"'), json_id_of(" ")]
                if self.state in ("OBJECT_COMMA",):
                    # either comma or close
                    return [json_id_of(","), json_id_of("}"), json_id_of(" ")]
                return [json_id_of('"'), json_id_of(" ")]
            elif container == "array":
                if self.state in ("ARRAY_START",):
                    # first value or close — no comma allowed
                    return self._value_tokens() + [json_id_of("]"), json_id_of(" ")]
                if self.state in ("ARRAY_AFTER_COMMA",):
                    # must begin a value — no close here
                    return self._value_tokens() + [json_id_of(" ")]
                if self.state in ("ARRAY_COMMA",):
                    # either comma or close
                    return [json_id_of(","), json_id_of("]"), json_id_of(" ")]
                return self._value_tokens() + [json_id_of(" ")]

        # Fallback: allow openers only (no closers when nothing to close)
        return [json_id_of("{"), json_id_of("[")]

    def reset(self):
        """Reset parser state for new generation."""
        self.state = "START"
        self.stack = []
        self.in_string = False
        self.string_escaped = False
        self.hex_count = 0
        self.hex_buffer = []
        self.depth = 0
        self.num_state = "NONE"
        self.ws_run = 0
        self.sur_hi_pending = False
        self.num_leading_zero = False
