llama.cpp/llama_cpp/llama_grammar.py

"""Python implementation of llama grammar parser directly translated from C++ source file in vendor/llama.cpp/common/grammar-parser.cpp."""

# flake8: noqa
from pathlib import Path
import sys
from ctypes import *  # type: ignore
from enum import Enum
from itertools import islice
from typing import (
    Callable,
    Generic,
    List,
    Optional,
    OrderedDict,
    TextIO,
    Tuple,
    TypeVar,
    Union,
    overload,
)

import llama_cpp.llama_cpp as llama_cpp

# Type aliases
llama_grammar_element = llama_cpp.llama_grammar_element
llama_grammar_element_p = llama_cpp.llama_grammar_element_p
llama_grammar_p = llama_cpp.llama_grammar_p

# Type variables
Ptr = TypeVar("Ptr", bound="const_char_p")
T = TypeVar("T")
U = TypeVar("U")
V = TypeVar("V")
W = TypeVar("W")


class Sentinel:
    """Used to mark the end of a iterator of std::vector & std::map."""


class LlamaGrammar:
    """Keeps reference counts of all the arguments, so that they are not
    garbage collected by Python."""

    def __del__(self) -> None:
        """Free the grammar pointer when the object is deleted."""
        if self.grammar is not None:
            llama_cpp.llama_grammar_free(self.grammar)
            self.grammar = None

    def __init__(
        self,
        parsed_grammar: "parse_state",
    ) -> None:
        """Initialize the grammar pointer from the parsed state."""
        self._grammar_rules = (
            parsed_grammar.c_rules()
        )  # type: std.vector[std.vector[LlamaGrammarElement]]
        self._n_rules = self._grammar_rules.size()  # type: int
        self._start_rule_index = parsed_grammar.symbol_ids.at("root")  # type: int
        self.grammar = self.init()

    @classmethod
    def from_string(cls, grammar: str, verbose: bool = True) -> "LlamaGrammar":
        parsed_grammar = parse(const_char_p(grammar))  # type: parse_state
        if parsed_grammar.rules.empty():
            raise ValueError(
                f"{cls.from_string.__name__}: error parsing grammar file: parsed_grammar.rules is empty"
            )
        if verbose:
            print(f"{cls.from_string.__name__} grammar:", file=sys.stderr)
            print_grammar(sys.stdout, parsed_grammar)
            print(file=sys.stderr)
        return cls(parsed_grammar)

    @classmethod
    def from_file(cls, file: Union[str, Path], verbose: bool = True) -> "LlamaGrammar":
        try:
            with open(file) as f:
                grammar = f.read()
        except Exception as err:
            raise Exception(
                f"{cls.from_file.__name__}: error reading grammar file: {err}"
            )

        if grammar:
            return cls.from_string(grammar, verbose=verbose)

        raise ValueError(
            f"{cls.from_file.__name__}: error parsing grammar file: params_grammer is empty"
        )

    def init(self) -> None:
        # Step 1: Convert LlamaGrammarElement to llama_grammar_element
        self._element_lists = [
            [
                llama_grammar_element(c_int(elem.type.value), c_uint32(elem.value))
                for elem in subvector
            ]
            for subvector in self._grammar_rules
        ]  # type: List[List[llama_grammar_element]]

        # Step 2: Convert each list to llama_grammar_element array and get pointer
        self._element_arrays = [
            (llama_grammar_element * len(sublist))(*sublist)
            for sublist in self._element_lists
        ]  # type: List[Array[llama_grammar_element]]

        # Step 3: Get pointer of each array
        self._element_array_pointers = [
            cast(subarray, llama_grammar_element_p) for subarray in self._element_arrays
        ]  # type: List[llama_grammar_element_p]

        # Step 4: Make array of these pointers and get its pointer
        self._rules = (llama_grammar_element_p * len(self._element_array_pointers))(
            *self._element_array_pointers
        )
        self.grammar = llama_cpp.llama_grammar_init(
            self._rules, c_size_t(self._n_rules), c_size_t(self._start_rule_index)
        )

    def reset(self) -> None:
        if self.grammar is not None:
            llama_cpp.llama_grammar_free(self.grammar)
        self.init()


class LlamaGrammarElement:
    def __init__(self, type: "llama_gretype", value: int):
        self.type = type
        self.value = value  # Unicode code point or rule ID


class const_char_p:
    """C++ implementation of const char *."""

    def __init__(self, value: Union[str, Ptr], move: Optional[int] = None):
        if isinstance(value, const_char_p):
            # We're copying an existing const_char_p
            self.value = value.value
            self.pos = value.pos + (move or 0)
            return

        # We're creating a new const_char_p
        self.value = value
        self.pos = move or 0

    def __str__(self) -> str:
        assert self.value is not None, "null pointer"
        return self.value[self.pos :]

    def __getitem__(self, index: int) -> str:
        value = str(self)
        return value[index] if index < len(value) else ""

    @overload
    def __add__(self: Ptr, other: int) -> Ptr:
        ...

    @overload
    def __add__(self: Ptr, other: Ptr) -> int:
        ...

    def __add__(self: Ptr, other: Union[int, Ptr]) -> Union[int, Ptr]:
        return (
            self.__class__(self.value, self.pos + other)
            if isinstance(other, int)
            else self.pos + other.pos
        )

    @overload
    def __sub__(self: Ptr, other: int) -> Ptr:
        ...

    @overload
    def __sub__(self: Ptr, other: Ptr) -> int:
        ...

    def __sub__(self: Ptr, other: Union[int, Ptr]) -> Union[int, Ptr]:
        return (
            self.__class__(self.value, self.pos - other)
            if isinstance(other, int)
            else self.pos - other.pos
        )

    def __eq__(self: Ptr, other: Ptr) -> bool:
        assert self.value == other.value, "comparing pointers from different strings"
        return self.pos == other.pos

    def __lt__(self: Ptr, other: Ptr) -> bool:
        assert self.value == other.value, "comparing pointers from different strings"
        return self.pos < other.pos

    def __gt__(self: Ptr, other: Ptr) -> bool:
        assert self.value == other.value, "comparing pointers from different strings"
        return self.pos > other.pos


class std:
    @staticmethod
    def string(ptr: const_char_p, length: Optional[int] = None) -> str:
        """C++ implementation of std::string constructor."""
        value = str(ptr)
        if length is not None:
            value = value[:length]
        return value

    class vector(Generic[T], List[T]):
        """C++ implementation of std::vector."""

        class iterator:
            def __init__(self, vector: "std.vector[T]", index: int):
                self._vector = vector
                self._index = index
                self._version = vector._version

            def _check_version(self):
                if self._version != self._vector._version:
                    raise RuntimeError("Iterator used after vector was modified.")

            def __iter__(self):
                return self

            def __next__(self) -> T:
                self._check_version()
                if self._index >= self._vector.size():
                    raise StopIteration
                value = self._vector[self._index]
                self._index += 1
                return value

            def __add__(self, value: int) -> "std.vector[T].iterator":
                return self.__class__(self._vector, self._index + value)

            def __sub__(self, value: int) -> "std.vector[T].iterator":
                return self.__class__(self._vector, self._index - value)

        def __init__(self):
            self._version = 0

        def modify(self):
            # This is a bit of a hack to make sure iterators are invalidated
            self._version += 1

        def push_back(self, value: T) -> None:
            self.modify()
            self.append(value)

        def pop_back(self) -> None:
            self.modify()
            if not self.empty():
                self.pop()

        def back(self) -> T:
            return self[-1]

        def size(self) -> int:
            return len(self)

        def clear(self) -> None:
            self.modify()
            super().clear()

        def empty(self) -> bool:
            return self.size() == 0

        def data(self) -> "std.vector[T]":
            return self

        def resize(
            self,
            new_size: int,
            fill_value_factory: Optional[Callable[[], T]] = None,
        ) -> None:
            if new_size > self.size():
                if fill_value_factory is None:
                    raise ValueError("A fill value factory function must be provided.")
                self.reserve(new_size, fill_value_factory)
            elif new_size < self.size():
                self[:] = self[:new_size]

        def reserve(self, capacity: int, fill_value_factory: Callable[[], T]) -> None:
            if capacity > self.size():
                fill_value = fill_value_factory()
                self.extend([fill_value] * (capacity - self.size()))

        def front(self) -> T:
            if not self.empty():
                return self[0]
            else:
                raise IndexError("Vector is empty.")

        def assign(self, count: int, value: T) -> None:
            self.clear()
            self.extend([value] * count)

        def insert(
            self,
            pos: "std.vector[T].iterator",
            first: "std.vector[T].iterator",
            last: "std.vector[T].iterator",
        ) -> None:
            self[pos._index : pos._index] = list(
                islice(first._vector, first._index, last._index)
            )

        def begin(self) -> "std.vector[T].iterator":
            return self.iterator(self, 0)

        def end(self) -> "std.vector[T].iterator":
            return self.iterator(self, self.size())

    class map(Generic[T, U], OrderedDict[T, U]):
        """C++ implementation of std::map."""

        class iterator(Generic[V, W]):
            def __init__(self, _map: "std.map[T, U]", key: Union[T, Sentinel]):
                self._map = _map
                self.iter = iter(_map)
                self.key = key
                self._advance()

            def _sanitize_key(self) -> T:
                if isinstance(self.key, Sentinel):
                    raise StopIteration
                return self.key

            def _advance(self) -> None:
                try:
                    while next(self.iter) != self.key:
                        pass
                except StopIteration:
                    self.key = Sentinel()

            def __next__(self) -> Tuple[T, U]:
                key = self._sanitize_key()
                if key in self._map:
                    value = self._map[key]
                    self._advance()
                    return key, value
                else:
                    raise StopIteration

            def get(self) -> Tuple[T, U]:
                key = self._sanitize_key()
                return key, self._map[key]

            @property
            def first(self) -> T:
                return self._sanitize_key()

            @property
            def second(self) -> U:
                return self._map[self._sanitize_key()]

        def insert(
            self, key: T, value: U
        ) -> Tuple["std.map[T, U].iterator[T, U]", bool]:
            if key in self:
                return self.iterator(self, key), False
            else:
                self[key] = value
                return self.iterator(self, key), True

        def find(self, key: T) -> "std.map[T, U].iterator[T, U]":
            if key in self:
                return self.iterator(self, key)
            else:
                return self.end()

        def at(self, key: T) -> U:
            if key in self:
                return self[key]
            else:
                raise KeyError("The provided key is not found in the map.")

        def erase(self, iterator: "std.map[T, U].iterator[T, U]") -> None:
            key = iterator.first
            if key in self:
                del self[key]

        def size(self) -> int:
            return len(self)

        def empty(self) -> bool:
            return self.size() == 0

        def lower_bound(self, key: T) -> "std.map[T, U].iterator[T, U]":
            try:
                keys = sorted(list(self.keys()))  # type: ignore
                for k in keys:
                    if k >= key:
                        return self.iterator(self, k)
                raise ValueError("No key found that is not less than the input key")
            except TypeError:
                raise TypeError("Keys of type T cannot be sorted.")

        def begin(self) -> "std.map[T, U].iterator[T, U]":
            return self.iterator(self, next(iter(self)))

        def end(self) -> "std.map[T, U].iterator[T, U]":
            return self.iterator(self, Sentinel())


# // grammar element type
# enum llama_gretype {
#     // end of rule definition
#     LLAMA_GRETYPE_END            = 0,

#     // start of alternate definition for rule
#     LLAMA_GRETYPE_ALT            = 1,

#     // non-terminal element: reference to rule
#     LLAMA_GRETYPE_RULE_REF       = 2,

#     // terminal element: character (code point)
#     LLAMA_GRETYPE_CHAR           = 3,

#     // inverse char(s) ([^a], [^a-b] [^abc])
#     LLAMA_GRETYPE_CHAR_NOT       = 4,

#     // modifies a preceding LLAMA_GRETYPE_CHAR or LLAMA_GRETYPE_CHAR_ALT to
#     // be an inclusive range ([a-z])
#     LLAMA_GRETYPE_CHAR_RNG_UPPER = 5,


#     // modifies a preceding LLAMA_GRETYPE_CHAR or
#     // LLAMA_GRETYPE_CHAR_RNG_UPPER to add an alternate char to match ([ab], [a-zA])
#     LLAMA_GRETYPE_CHAR_ALT       = 6,
# };
class llama_gretype(Enum):
    """grammar element type"""

    LLAMA_GRETYPE_END = 0  # end of rule definition
    LLAMA_GRETYPE_ALT = 1  # start of alternate definition for rule
    LLAMA_GRETYPE_RULE_REF = 2  # non-terminal element: reference to rule
    LLAMA_GRETYPE_CHAR = 3  # terminal element: character (code point)
    LLAMA_GRETYPE_CHAR_NOT = 4  # inverse char(s) ([^a], [^a-b] [^abc])
    LLAMA_GRETYPE_CHAR_RNG_UPPER = 5  # modifies a preceding LLAMA_GRETYPE_CHAR or LLAMA_GRETYPE_CHAR_ALT to be an inclusive range ([a-z])
    LLAMA_GRETYPE_CHAR_ALT = 6  # modifies a preceding LLAMA_GRETYPE_CHAR or LLAMA_GRETYPE_CHAR_RNG_UPPER to add an alternate char to match ([ab], [a-zA])


# struct parse_state {
#     std::map<std::string, uint32_t>                 symbol_ids;
#     std::vector<std::vector<llama_grammar_element>> rules;
#     std::vector<const llama_grammar_element *> c_rules();
# };
class parse_state:
    def __init__(self):
        self.symbol_ids: std.map[str, int] = std.map()
        self.rules: std.vector[std.vector[LlamaGrammarElement]] = std.vector()

    # std::vector<const llama_grammar_element *> parse_state::c_rules() {
    #     std::vector<const llama_grammar_element *> ret;
    #     for (const auto & rule : rules) {
    #         ret.push_back(rule.data());
    #     }
    #     return ret;
    # }
    def c_rules(self) -> std.vector[std.vector[LlamaGrammarElement]]:
        ret = std.vector()  # type: std.vector[std.vector[LlamaGrammarElement]]
        for rule in self.rules:
            ret.push_back(rule.data())
        return ret

    def __repr__(self) -> str:
        return (
            f"parse_state(symbol_ids={len(self.symbol_ids)}, rules={len(self.rules)})"
        )


# struct llama_grammar {
#     const std::vector<std::vector<llama_grammar_element>>   rules;
#     std::vector<std::vector<const llama_grammar_element *>> stacks;
# };
# class llama_grammar:
#     def __init__(
#         self,
#         rules: std.vector[std.vector[llama_grammar_element]],
#         stacks: std.vector[std.vector[llama_grammar_element]],
#     ):
#         self.rules = rules
#         self.stacks = stacks


# uint32_t get_symbol_id(parse_state & state, const char * src, size_t len) {
#     uint32_t next_id = static_cast<uint32_t>(state.symbol_ids.size());
#     auto result = state.symbol_ids.insert(std::make_pair(std::string(src, len), next_id));
#     return result.first->second;
# }
def get_symbol_id(state: parse_state, src: const_char_p, len: int) -> int:
    next_id = state.symbol_ids.size()  # type: int
    result = state.symbol_ids.insert(std.string(src, len), next_id)
    return result[0].second  # type: ignore


# uint32_t generate_symbol_id(parse_state & state, const std::string & base_name) {
#     uint32_t next_id = static_cast<uint32_t>(state.symbol_ids.size());
#     state.symbol_ids[base_name + '_' + std::to_string(next_id)] = next_id;
#     return next_id;
# }
def generate_symbol_id(state: parse_state, base_name: str) -> int:
    next_id = state.symbol_ids.size()  # type: int
    state.symbol_ids[base_name + "_" + str(next_id)] = next_id
    return next_id


# void add_rule(
#         parse_state & state,
#         uint32_t      rule_id,
#         const std::vector<llama_grammar_element> & rule) {
#     if (state.rules.size() <= rule_id) {
#         state.rules.resize(rule_id + 1);
#     }
#     state.rules[rule_id] = rule;
# }
def add_rule(
    state: parse_state,
    rule_id: int,
    rule: std.vector[LlamaGrammarElement],
) -> None:
    if state.rules.size() <= rule_id:
        state.rules.resize(
            rule_id + 1,
            fill_value_factory=std.vector[LlamaGrammarElement],
        )
    state.rules[rule_id] = rule


# std::pair<uint32_t, const char *> decode_utf8(const char * src) {
#     static const int lookup[] = { 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 3, 4 };
#     uint8_t  first_byte = static_cast<uint8_t>(*src);
#     uint8_t  highbits   = first_byte >> 4;
#     int      len        = lookup[highbits];
#     uint8_t  mask       = (1 << (8 - len)) - 1;
#     uint32_t value      = first_byte & mask;
#     const char * end    = src + len; // may overrun!
#     const char * pos    = src + 1;
#     for ( ; pos < end && *pos; pos++) {
#         value = (value << 6) + (static_cast<uint8_t>(*pos) & 0x3F);
#     }
#     return std::make_pair(value, pos);
# }
def decode_utf8(src: const_char_p) -> Tuple[int, const_char_p]:
    """Decodes a UTF-8 character from the source string."""
    lookup = (1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 3, 4)
    first_byte = ord(src[0])  # type: int
    highbits = first_byte >> 4  # type: int
    len = lookup[highbits]  # type: int
    mask = (1 << (8 - len)) - 1  # type: int
    value = first_byte & mask  # type: int
    end = src + len  # type: const_char_p # may overrun!
    pos = src + 1  # type: const_char_p
    while pos < end and pos[0]:
        value = (value << 6) + (ord(pos[0]) & 0x3F)
        pos += 1
    return value, pos


# bool is_word_char(char c) {
#     return ('a' <= c && c <= 'z') || ('A' <= c && c <= 'Z') || c == '-' || ('0' <= c && c <= '9');
# }
def is_word_char(c: str) -> bool:
    return ("a" <= c <= "z") or ("A" <= c <= "Z") or c == "-" or ("0" <= c <= "9")


# std::pair<uint32_t, const char *> parse_hex(const char * src, int size) {
#     const char * pos   = src;
#     const char * end   = src + size;
#     uint32_t     value = 0;
#     for ( ; pos < end && *pos; pos++) {
#         value <<= 4;
#         char c = *pos;
#         if ('a' <= c && c <= 'f') {
#             value += c - 'a' + 10;
#         } else if ('A' <= c && c <= 'F') {
#             value += c - 'A' + 10;
#         } else if ('0' <= c && c <= '9') {
#             value += c - '0';
#         } else {
#             break;
#         }
#     }
#     if (pos != end) {
#         throw std::runtime_error("expecting " + std::to_string(size) + " hex chars at " + src);
#     }
#     return std::make_pair(value, pos);
# }
def parse_hex(src: const_char_p, size: int) -> Tuple[int, const_char_p]:
    pos = const_char_p(src)  # type: const_char_p
    end = src + size  # type: const_char_p
    value = 0  # type: int
    while pos < end and pos[0]:
        value <<= 4
        c = pos[0]  # type: str
        if "a" <= c <= "f":
            value += ord(c) - ord("a") + 10
        elif "A" <= c <= "F":
            value += ord(c) - ord("A") + 10
        elif "0" <= c <= "9":
            value += ord(c) - ord("0")
        else:
            break
        pos += 1
    if pos != end:
        raise RuntimeError("expecting " + str(size) + " hex chars at " + str(src))
    return (value, pos)


# std::pair<uint32_t, const char *> parse_char(const char * src) {
#     if (*src == '\\') {
#         switch (src[1]) {
#             case 'x': return parse_hex(src + 2, 2);
#             case 'u': return parse_hex(src + 2, 4);
#             case 'U': return parse_hex(src + 2, 8);
#             case 't': return std::make_pair('\t', src + 2);
#             case 'r': return std::make_pair('\r', src + 2);
#             case 'n': return std::make_pair('\n', src + 2);
#             case '\\':
#             case '"':
#             case '[':
#             case ']':
#                 return std::make_pair(src[1], src + 2);
#             default:
#                 throw std::runtime_error(std::string("unknown escape at ") + src);
#         }
#     } else if (*src) {
#         return decode_utf8(src);
#     }
#     throw std::runtime_error("unexpected end of input");
# }
def parse_char(src: const_char_p) -> Tuple[int, const_char_p]:
    if src[0] == "\\":
        case = src[1]  # type: str
        if case == "x":
            return parse_hex(src + 2, 2)
        elif case == "u":
            return parse_hex(src + 2, 4)
        elif case == "U":
            return parse_hex(src + 2, 8)
        elif case == "t":
            return (ord("\t"), src + 2)  # implicit cast
        elif case == "r":
            return (ord("\r"), src + 2)  # implicit cast
        elif case == "n":
            return (ord("\n"), src + 2)  # implicit cast
        elif case in ("\\", '"', "[", "]"):
            return (ord(case), src + 2)  # implicit cast
        else:
            raise RuntimeError("unknown escape at " + str(src))
    elif src[0]:
        return decode_utf8(src)
    else:
        raise RuntimeError("unexpected end of input")


# const char * parse_name(const char * src) {
#     const char * pos = src;
#     while (is_word_char(*pos)) {
#         pos++;
#     }
#     if (pos == src) {
#         throw std::runtime_error(std::string("expecting name at ") + src);
#     }
#     return pos;
# }
def parse_name(src: const_char_p) -> const_char_p:
    pos = const_char_p(src)  # type: const_char_p
    while is_word_char(pos[0]):
        pos += 1
    if pos == src:
        raise RuntimeError("expecting name at " + str(src))
    return pos


# const char * parse_space(const char * src, bool newline_ok) {
#     const char * pos = src;
#     while (*pos == ' ' || *pos == '\t' || *pos == '#' ||
#             (newline_ok && (*pos == '\r' || *pos == '\n'))) {
#         if (*pos == '#') {
#             while (*pos && *pos != '\r' && *pos != '\n') {
#                 pos++;
#             }
#         } else {
#             pos++;
#         }
#     }
#     return pos;
# }
def parse_space(src: const_char_p, newline_ok: bool) -> const_char_p:
    pos = const_char_p(src)  # type: const_char_p
    while pos[0] in (" ", "\t", "#") or (newline_ok and pos[0] in ("\r", "\n")):
        if pos[0] == "#":
            while pos[0] is not None and pos[0] not in ("\r", "\n"):
                pos += 1
        else:
            pos += 1
    return pos


# const char * parse_sequence(
#         parse_state                        & state,
#         const char                         * src,
#         const std::string                  & rule_name,
#         std::vector<llama_grammar_element> & out_elements,
#         bool                                 is_nested) {
def parse_sequence(
    state: parse_state,
    src: const_char_p,
    rule_name: str,
    out_elements: std.vector[LlamaGrammarElement],
    is_nested: bool,
) -> const_char_p:
    # size_t last_sym_start = out_elements.size();
    # const char * pos = src;
    last_sym_start = out_elements.size()  # type: int
    pos = const_char_p(src)  # type: const_char_p
    # while (*pos) {
    while pos[0]:
        # if (*pos == '"') { // literal string
        #     pos++;
        #     last_sym_start = out_elements.size();
        #     while (*pos != '"') {
        #         auto char_pair = parse_char(pos);
        #                 pos       = char_pair.second;
        #         out_elements.push_back({LLAMA_GRETYPE_CHAR, char_pair.first});
        #     }
        #     pos = parse_space(pos + 1, is_nested);
        if pos[0] == '"':  # literal string
            pos += 1
            last_sym_start = out_elements.size()
            while pos[0] != '"':
                char_pair = parse_char(pos)  # type: Tuple[int, const_char_p]
                pos = char_pair[1]
                out_elements.push_back(
                    LlamaGrammarElement(llama_gretype.LLAMA_GRETYPE_CHAR, char_pair[0])
                )
            pos = parse_space(pos + 1, is_nested)
        # } else if (*pos == '[') { // char range(s)
        #     pos++;
        #     enum llama_gretype start_type = LLAMA_GRETYPE_CHAR;
        elif pos[0] == "[":  # char range(s)
            pos += 1
            start_type = llama_gretype.LLAMA_GRETYPE_CHAR  # type: llama_gretype
            # if (*pos == '^') {
            #     pos++;
            #     start_type = LLAMA_GRETYPE_CHAR_NOT;
            # }
            # last_sym_start = out_elements.size();
            if pos[0] == "^":
                pos += 1
                start_type = llama_gretype.LLAMA_GRETYPE_CHAR_NOT
            last_sym_start = out_elements.size()
            # while (*pos != ']') {
            #     auto char_pair = parse_char(pos);
            #             pos       = char_pair.second;
            #     enum llama_gretype type = last_sym_start < out_elements.size()
            #         ? LLAMA_GRETYPE_CHAR_ALT
            #         : start_type;
            #     out_elements.push_back({type, char_pair.first});
            while pos[0] != "]":
                char_pair = parse_char(pos)  # type: Tuple[int, const_char_p]
                pos = char_pair[1]
                type = (
                    llama_gretype.LLAMA_GRETYPE_CHAR_ALT
                    if last_sym_start < out_elements.size()
                    else start_type
                )  # type: llama_gretype
                out_elements.push_back(LlamaGrammarElement(type, char_pair[0]))
                #     if (pos[0] == '-' && pos[1] != ']') {
                #         auto endchar_pair = parse_char(pos + 1);
                #                 pos          = endchar_pair.second;
                #         out_elements.push_back({LLAMA_GRETYPE_CHAR_RNG_UPPER, endchar_pair.first});
                #     }
                # }
                if pos[0] == "-" and pos[1] != "]":
                    endchar_pair = parse_char(pos + 1)  # type: Tuple[int, const_char_p]
                    pos = endchar_pair[1]
                    out_elements.push_back(
                        LlamaGrammarElement(
                            llama_gretype.LLAMA_GRETYPE_CHAR_RNG_UPPER,
                            endchar_pair[0],
                        )
                    )
            # pos = parse_space(pos + 1, is_nested);
            pos = parse_space(pos + 1, is_nested)
        # } else if (is_word_char(*pos)) { // rule reference
        #     const char * name_end    = parse_name(pos);
        #     uint32_t     ref_rule_id = get_symbol_id(state, pos, name_end - pos);
        #     pos = parse_space(name_end, is_nested);
        #     last_sym_start = out_elements.size();
        #     out_elements.push_back({LLAMA_GRETYPE_RULE_REF, ref_rule_id});
        elif is_word_char(pos[0]):  # rule reference
            name_end = parse_name(pos)  # type: const_char_p
            ref_rule_id = get_symbol_id(state, pos, name_end - pos)  # type: int
            pos = parse_space(name_end, is_nested)
            last_sym_start = out_elements.size()
            out_elements.push_back(
                LlamaGrammarElement(llama_gretype.LLAMA_GRETYPE_RULE_REF, ref_rule_id)
            )
        # } else if (*pos == '(') { // grouping
        #     // parse nested alternates into synthesized rule
        #     pos = parse_space(pos + 1, true);
        #     uint32_t sub_rule_id = generate_symbol_id(state, rule_name);
        #     pos = parse_alternates(state, pos, rule_name, sub_rule_id, true);
        #     last_sym_start = out_elements.size();
        #     // output reference to synthesized rule
        #     out_elements.push_back({LLAMA_GRETYPE_RULE_REF, sub_rule_id});
        #     if (*pos != ')') {
        #         throw std::runtime_error(std::string("expecting ')' at ") + pos);
        #     }
        #     pos = parse_space(pos + 1, is_nested);
        elif pos[0] == "(":  # grouping
            # parse nested alternates into synthesized rule
            pos = parse_space(pos + 1, True)
            sub_rule_id = generate_symbol_id(state, rule_name)  # type: int
            pos = parse_alternates(state, pos, rule_name, sub_rule_id, True)
            last_sym_start = out_elements.size()
            # output reference to synthesized rule
            out_elements.push_back(
                LlamaGrammarElement(llama_gretype.LLAMA_GRETYPE_RULE_REF, sub_rule_id)
            )
            if pos[0] != ")":
                raise RuntimeError("expecting ')' at " + str(pos))
            pos = parse_space(pos + 1, is_nested)
        # } else if (*pos == '*' || *pos == '+' || *pos == '?') { // repetition operator
        #     if (last_sym_start == out_elements.size()) {
        #         throw std::runtime_error(std::string("expecting preceeding item to */+/? at ") + pos);
        #     }
        elif pos[0] in ("*", "+", "?"):  # repetition operator
            if last_sym_start == out_elements.size():
                raise RuntimeError("expecting preceding item to */+/? at " + str(pos))
            # // apply transformation to previous symbol (last_sym_start to end) according to
            # // rewrite rules:
            # // S* --> S' ::= S S' |
            # // S+ --> S' ::= S S' | S
            # // S? --> S' ::= S |
            # uint32_t sub_rule_id = generate_symbol_id(state, rule_name);
            # std::vector<llama_grammar_element> sub_rule;
            # // add preceding symbol to generated rule
            # sub_rule.insert(
            #     sub_rule.end(), out_elements.begin() + last_sym_start, out_elements.end());
            sub_rule_id = generate_symbol_id(state, rule_name)  # type: int
            sub_rule = std.vector[
                LlamaGrammarElement
            ]()  # type: std.vector[LlamaGrammarElement]
            sub_rule.insert(
                sub_rule.end(),
                out_elements.begin() + last_sym_start,
                out_elements.end(),
            )
            # if (*pos == '*' || *pos == '+') {
            #     // cause generated rule to recurse
            #     sub_rule.push_back({LLAMA_GRETYPE_RULE_REF, sub_rule_id});
            # }
            # // mark start of alternate def
            # sub_rule.push_back({LLAMA_GRETYPE_ALT, 0});
            if pos[0] in ("*", "+"):
                sub_rule.push_back(
                    LlamaGrammarElement(
                        llama_gretype.LLAMA_GRETYPE_RULE_REF, sub_rule_id
                    )
                )
            sub_rule.push_back(LlamaGrammarElement(llama_gretype.LLAMA_GRETYPE_ALT, 0))
            # if (*pos == '+') {
            #     // add preceding symbol as alternate only for '+' (otherwise empty)
            #     sub_rule.insert(
            #         sub_rule.end(), out_elements.begin() + last_sym_start, out_elements.end());
            # }
            # sub_rule.push_back({LLAMA_GRETYPE_END, 0});
            # add_rule(state, sub_rule_id, sub_rule);
            # // in original rule, replace previous symbol with reference to generated rule
            # out_elements.resize(last_sym_start);
            # out_elements.push_back({LLAMA_GRETYPE_RULE_REF, sub_rule_id});
            # pos = parse_space(pos + 1, is_nested);
            if pos[0] == "+":
                # add preceding symbol as alternate only for '+' (otherwise empty)
                sub_rule.insert(
                    sub_rule.end(),
                    out_elements.begin() + last_sym_start,
                    out_elements.end(),
                )
            sub_rule.push_back(LlamaGrammarElement(llama_gretype.LLAMA_GRETYPE_END, 0))
            add_rule(state, sub_rule_id, sub_rule)
            # in original rule, replace previous symbol with reference to generated rule
            out_elements.resize(last_sym_start)
            out_elements.push_back(
                LlamaGrammarElement(llama_gretype.LLAMA_GRETYPE_RULE_REF, sub_rule_id)
            )
            pos = parse_space(pos + 1, is_nested)
        # } else {
        #     break;
        # }
        else:
            break
    #     }
    #     return pos;
    # }
    return pos


# const char * parse_alternates(
#         parse_state       & state,
#         const char        * src,
#         const std::string & rule_name,
#         uint32_t            rule_id,
#         bool                is_nested) {
#     std::vector<llama_grammar_element> rule;
#     const char * pos = parse_sequence(state, src, rule_name, rule, is_nested);
#     while (*pos == '|') {
#         rule.push_back({LLAMA_GRETYPE_ALT, 0});
#         pos = parse_space(pos + 1, true);
#         pos = parse_sequence(state, pos, rule_name, rule, is_nested);
#     }
#     rule.push_back({LLAMA_GRETYPE_END, 0});
#     add_rule(state, rule_id, rule);
#     return pos;
# }
def parse_alternates(
    state: parse_state,
    src: const_char_p,
    rule_name: str,
    rule_id: int,
    is_nested: bool,
) -> const_char_p:
    rule = std.vector()  # type: std.vector[LlamaGrammarElement]
    pos = parse_sequence(state, src, rule_name, rule, is_nested)  # type: const_char_p
    while pos[0] == "|":
        rule.push_back(LlamaGrammarElement(llama_gretype.LLAMA_GRETYPE_ALT, 0))
        pos = parse_space(pos + 1, True)
        pos = parse_sequence(state, pos, rule_name, rule, is_nested)
    rule.push_back(LlamaGrammarElement(llama_gretype.LLAMA_GRETYPE_END, 0))
    add_rule(state, rule_id, rule)
    return pos


# const char * parse_rule(parse_state & state, const char * src) {
#     const char * name_end = parse_name(src);
#     const char * pos      = parse_space(name_end, false);
#     size_t       name_len = name_end - src;
#     uint32_t     rule_id  = get_symbol_id(state, src, name_len);
#     const std::string name(src, name_len);

#     if (!(pos[0] == ':' && pos[1] == ':' && pos[2] == '=')) {
#         throw std::runtime_error(std::string("expecting ::= at ") + pos);
#     }
#     pos = parse_space(pos + 3, true);

#     pos = parse_alternates(state, pos, name, rule_id, false);


#     if (*pos == '\r') {
#         pos += pos[1] == '\n' ? 2 : 1;
#     } else if (*pos == '\n') {
#         pos++;
#     } else if (*pos) {
#         throw std::runtime_error(std::string("expecting newline or end at ") + pos);
#     }
#     return parse_space(pos, true);
# }
def parse_rule(state: parse_state, src: const_char_p) -> const_char_p:
    name_end = parse_name(src)  # type: const_char_p
    pos = parse_space(name_end, False)  # type: const_char_p
    name_len = name_end - src  # type: int
    rule_id = get_symbol_id(state, src, name_len)  # type: int
    name = std.string(src, name_len)  # type: str

    if not (pos[0] == ":" and pos[1] == ":" and pos[2] == "="):
        raise RuntimeError("expecting ::= at " + str(pos))

    pos = parse_space(pos + 3, True)  # type: const_char_p
    pos = parse_alternates(state, pos, name, rule_id, False)  # type: const_char_p

    if pos[0] == "\r":
        pos += 2 if pos[1] == "\n" else 1
    elif pos[0] == "\n":
        pos += 1
    elif pos[0]:
        raise RuntimeError("expecting newline or end at " + str(pos))
    return parse_space(pos, True)


# parse_state parse(const char * src) {
#     try {
#         parse_state state;
#         const char * pos = parse_space(src, true);
#         while (*pos) {
#             pos = parse_rule(state, pos);
#         }
#         return state;
#     } catch (const std::exception & err) {
#         fprintf(stderr, "%s: error parsing grammar: %s\n", __func__, err.what());
#         return parse_state();
#     }
# }
def parse(src: const_char_p) -> parse_state:
    try:
        state = parse_state()  # type: parse_state
        pos = parse_space(src, True)  # type: const_char_p
        while pos[0]:
            pos = parse_rule(state, pos)
        return state
    except Exception as err:
        print(f"{parse.__name__}: error parsing grammar: {err}")
        return parse_state()


# void print_grammar_char(FILE * file, uint32_t c) {
#     if (0x20 <= c && c <= 0x7f) {
#         fprintf(file, "%c", static_cast<char>(c));
#     } else {
#         // cop out of encoding UTF-8
#         fprintf(file, "<U+%04X>", c);
#     }
# }
def print_grammar_char(file: TextIO, c: int) -> None:
    if 0x20 <= c and c <= 0x7F:
        file.write(chr(c))
    else:
        # cop out of encoding UTF-8
        file.write(f"<U+{c:04X}>")


# bool is_char_element(llama_grammar_element elem) {
#     switch (elem.type) {
#         case LLAMA_GRETYPE_CHAR:           return true;
#         case LLAMA_GRETYPE_CHAR_NOT:       return true;
#         case LLAMA_GRETYPE_CHAR_ALT:       return true;
#         case LLAMA_GRETYPE_CHAR_RNG_UPPER: return true;
#         default:                           return false;
#     }
# }
def is_char_element(elem: LlamaGrammarElement) -> bool:
    return elem.type in (
        llama_gretype.LLAMA_GRETYPE_CHAR,
        llama_gretype.LLAMA_GRETYPE_CHAR_NOT,
        llama_gretype.LLAMA_GRETYPE_CHAR_ALT,
        llama_gretype.LLAMA_GRETYPE_CHAR_RNG_UPPER,
    )


# void print_rule(
#         FILE     * file,
#         uint32_t   rule_id,
#         const std::vector<llama_grammar_element> & rule,
#         const std::map<uint32_t, std::string>    & symbol_id_names) {
def print_rule(
    file: TextIO,
    rule_id: int,
    rule: std.vector[LlamaGrammarElement],
    symbol_id_names: std.map[int, str],
) -> None:
    #     if (rule.empty() || rule.back().type != LLAMA_GRETYPE_END) {
    #         throw std::runtime_error(
    #             "malformed rule, does not end with LLAMA_GRETYPE_END: " + std::to_string(rule_id));
    #     }
    #     fprintf(file, "%s ::= ", symbol_id_names.at(rule_id).c_str());
    if rule.empty() or rule.back().type != llama_gretype.LLAMA_GRETYPE_END:
        raise RuntimeError(
            "malformed rule, does not end with LLAMA_GRETYPE_END: " + str(rule_id)
        )
    print(f"{symbol_id_names.at(rule_id)} ::=", file=file, end=" ")
    #     for (size_t i = 0, end = rule.size() - 1; i < end; i++) {
    #         llama_grammar_element elem = rule[i];
    #         switch (elem.type) {
    #             case LLAMA_GRETYPE_END:
    #                 throw std::runtime_error(
    #                     "unexpected end of rule: " + std::to_string(rule_id) + "," +
    #                     std::to_string(i));
    #             case LLAMA_GRETYPE_ALT:
    #                 fprintf(file, "| ");
    #                 break;
    #             case LLAMA_GRETYPE_RULE_REF:
    #                 fprintf(file, "%s ", symbol_id_names.at(elem.value).c_str());
    #                 break;
    #             case LLAMA_GRETYPE_CHAR:
    #                 fprintf(file, "[");
    #                 print_grammar_char(file, elem.value);
    #                 break;
    #             case LLAMA_GRETYPE_CHAR_NOT:
    #                 fprintf(file, "[^");
    #                 print_grammar_char(file, elem.value);
    #                 break;
    #             case LLAMA_GRETYPE_CHAR_RNG_UPPER:
    #                 if (i == 0 || !is_char_element(rule[i - 1])) {
    #                     throw std::runtime_error(
    #                         "LLAMA_GRETYPE_CHAR_RNG_UPPER without preceding char: " +
    #                         std::to_string(rule_id) + "," + std::to_string(i));
    #                 }
    #                 fprintf(file, "-");
    #                 print_grammar_char(file, elem.value);
    #                 break;
    #             case LLAMA_GRETYPE_CHAR_ALT:
    #                 if (i == 0 || !is_char_element(rule[i - 1])) {
    #                     throw std::runtime_error(
    #                         "LLAMA_GRETYPE_CHAR_ALT without preceding char: " +
    #                         std::to_string(rule_id) + "," + std::to_string(i));
    #                 }
    #                 print_grammar_char(file, elem.value);
    #                 break;
    #         }
    for i, elem in enumerate(rule[:-1]):
        case = elem.type  # type: llama_gretype
        if case is llama_gretype.LLAMA_GRETYPE_END:
            raise RuntimeError("unexpected end of rule: " + str(rule_id) + "," + str(i))
        elif case is llama_gretype.LLAMA_GRETYPE_ALT:
            print("| ", file=file, end="")
        elif case is llama_gretype.LLAMA_GRETYPE_RULE_REF:
            print(f"{symbol_id_names.at(elem.value)} ", file=file, end="")
        elif case is llama_gretype.LLAMA_GRETYPE_CHAR:
            print("[", file=file, end="")
            print_grammar_char(file, elem.value)
        elif case is llama_gretype.LLAMA_GRETYPE_CHAR_NOT:
            print("[^", file=file, end="")
            print_grammar_char(file, elem.value)
        elif case is llama_gretype.LLAMA_GRETYPE_CHAR_RNG_UPPER:
            if i == 0 or not is_char_element(rule[i - 1]):
                raise RuntimeError(
                    "LLAMA_GRETYPE_CHAR_RNG_UPPER without preceding char: "
                    + str(rule_id)
                    + ","
                    + str(i)
                )
            print("-", file=file, end="")
            print_grammar_char(file, elem.value)
        elif case is llama_gretype.LLAMA_GRETYPE_CHAR_ALT:
            if i == 0 or not is_char_element(rule[i - 1]):
                raise RuntimeError(
                    "LLAMA_GRETYPE_CHAR_ALT without preceding char: "
                    + str(rule_id)
                    + ","
                    + str(i)
                )
            print_grammar_char(file, elem.value)
        # if (is_char_element(elem)) {
        #     switch (rule[i + 1].type) {
        #         case LLAMA_GRETYPE_CHAR_ALT:
        #         case LLAMA_GRETYPE_CHAR_RNG_UPPER:
        #             break;
        #         default:
        #             fprintf(file, "] ");
        if is_char_element(elem):
            if rule[i + 1].type in (
                llama_gretype.LLAMA_GRETYPE_CHAR_ALT,
                llama_gretype.LLAMA_GRETYPE_CHAR_RNG_UPPER,
            ):
                pass
            else:
                print("] ", file=file, end="")
    #             }
    #         }
    #     }
    #     fprintf(file, "\n");
    # }
    print(file=file)


# void print_grammar(FILE * file, const parse_state & state) {
#     try {
#         std::map<uint32_t, std::string> symbol_id_names;
#         for (auto kv : state.symbol_ids) {
#             symbol_id_names[kv.second] = kv.first;
#         }
#         for (size_t i = 0, end = state.rules.size(); i < end; i++) {
#             // fprintf(file, "%zu: ", i);
#             // print_rule_binary(file, state.rules[i]);
#             print_rule(file, i, state.rules[i], symbol_id_names);
#             // fprintf(file, "\n");
#         }
#     } catch (const std::exception & err) {
#         fprintf(stderr, "\n%s: error printing grammar: %s\n", __func__, err.what());
#     }
# }
def print_grammar(file: TextIO, state: parse_state) -> None:
    try:
        symbol_id_names = std.map()  # type: std.map[int, str]
        for kv in state.symbol_ids.items():
            symbol_id_names[kv[1]] = kv[0]

        for i, rule in enumerate(state.rules):
            print_rule(file, i, rule, symbol_id_names)
    except Exception as err:
        print(
            f"{print_grammar.__name__}: error printing grammar: {err}",
            file=sys.stderr,
        )


"""llama.cpp gbnf rules from vendor/llama.cpp/grammars"""

ARITHMETIC_GBNF = r"""
root  ::= (expr "=" ws term "\n")+
expr  ::= term ([-+*/] term)*
term  ::= ident | num | "(" ws expr ")" ws
ident ::= [a-z] [a-z0-9_]* ws
num   ::= [0-9]+ ws
ws    ::= [ \t\n]*
"""

C_GBNF = r"""
root ::= (declaration)*

declaration ::= dataType identifier "(" parameter? ")" "{" statement* "}"

dataType  ::= "int" ws | "float" ws | "char" ws
identifier ::= [a-zA-Z_] [a-zA-Z_0-9]*

parameter ::= dataType identifier

statement ::=
    ( dataType identifier ws "=" ws expression ";" ) |
    ( identifier ws "=" ws expression ";" ) |
    ( identifier ws "(" argList? ")" ";" ) |
    ( "return" ws expression ";" ) |
    ( "while" "(" condition ")" "{" statement* "}" ) |
    ( "for" "(" forInit ";" ws condition ";" ws forUpdate ")" "{" statement* "}" ) |
    ( "if" "(" condition ")" "{" statement* "}" ("else" "{" statement* "}")? ) |
    ( singleLineComment ) |
    ( multiLineComment )

forInit ::= dataType identifier ws "=" ws expression | identifier ws "=" ws expression
forUpdate ::= identifier ws "=" ws expression

condition ::= expression relationOperator expression
relationOperator ::= ("<=" | "<" | "==" | "!=" | ">=" | ">")

expression ::= term (("+" | "-") term)*
term ::= factor(("*" | "/") factor)*

factor ::= identifier | number | unaryTerm | funcCall | parenExpression
unaryTerm ::= "-" factor
funcCall ::= identifier "(" argList? ")"
parenExpression ::= "(" ws expression ws ")"

argList ::= expression ("," ws expression)*

number ::= [0-9]+

singleLineComment ::= "//" [^\n]* "\n"
multiLineComment ::= "/*" ( [^*] | ("*" [^/]) )* "*/"

ws ::= ([ \t\n]+)
"""

CHESS_GBNF = r"""
root   ::= object
value  ::= object | array | string | number | ("true" | "false" | "null") ws

object ::=
  "{" ws (
            string ":" ws value
    ("," ws string ":" ws value)*
  )? "}" ws

array  ::=
  "[" ws (
            value
    ("," ws value)*
  )? "]" ws

string ::=
  "\"" (
    [^"\\] |
    "\\" (["\\/bfnrt] | "u" [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F]) # escapes
  )* "\"" ws

number ::= ("-"? ([0-9] | [1-9] [0-9]*)) ("." [0-9]+)? ([eE] [-+]? [0-9]+)? ws

# Optional space: by convention, applied in this grammar after literal chars when allowed
ws ::= ([ \t\n] ws)?
"""

JAPANESE_GBNF = r"""
root   ::= object
value  ::= object | array | string | number | ("true" | "false" | "null") ws

object ::=
  "{" ws (
            string ":" ws value
    ("," ws string ":" ws value)*
  )? "}" ws

array  ::=
  "[" ws (
            value
    ("," ws value)*
  )? "]" ws

string ::=
  "\"" (
    [^"\\] |
    "\\" (["\\/bfnrt] | "u" [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F]) # escapes
  )* "\"" ws

number ::= ("-"? ([0-9] | [1-9] [0-9]*)) ("." [0-9]+)? ([eE] [-+]? [0-9]+)? ws

# Optional space: by convention, applied in this grammar after literal chars when allowed
ws ::= ([ \t\n] ws)?
"""

JSON_ARR_GBNF = r"""
# This is the same as json.gbnf but we restrict whitespaces at the end of the root array
# Useful for generating JSON arrays

root   ::= arr
value  ::= object | array | string | number | ("true" | "false" | "null") ws

arr  ::=
  "[\n" ws (
            value
    (",\n" ws value)*
  )? "]"

object ::=
  "{" ws (
            string ":" ws value
    ("," ws string ":" ws value)*
  )? "}" ws

array  ::=
  "[" ws (
            value
    ("," ws value)*
  )? "]" ws

string ::=
  "\"" (
    [^"\\] |
    "\\" (["\\/bfnrt] | "u" [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F]) # escapes
  )* "\"" ws

number ::= ("-"? ([0-9] | [1-9] [0-9]*)) ("." [0-9]+)? ([eE] [-+]? [0-9]+)? ws

# Optional space: by convention, applied in this grammar after literal chars when allowed
ws ::= ([ \t\n] ws)?
"""


JSON_GBNF = r"""
root   ::= object
value  ::= object | array | string | number | ("true" | "false" | "null") ws

object ::=
  "{" ws (
            string ":" ws value
    ("," ws string ":" ws value)*
  )? "}" ws

array  ::=
  "[" ws (
            value
    ("," ws value)*
  )? "]" ws

string ::=
  "\"" (
    [^"\\] |
    "\\" (["\\/bfnrt] | "u" [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F]) # escapes
  )* "\"" ws

number ::= ("-"? ([0-9] | [1-9] [0-9]*)) ("." [0-9]+)? ([eE] [-+]? [0-9]+)? ws

ws ::= ([ \t\n] ws)?
"""

LIST_GBNF = r"""
root ::= item+

# Excludes various line break characters
item ::= "- " [^\r\n\x0b\x0c\x85\u2028\u2029]+ "\n"
"""

"""llama.cpp json-schema to grammar converter from vendor/llama.cpp/examples/json-schema-to-grammar.py"""
import json
import re
from typing import List, Optional

# whitespace is constrained to a single space char to prevent model "running away" in
# whitespace. Also maybe improves generation quality?
SPACE_RULE = '" "?'

PRIMITIVE_RULES = {
    "boolean": '("true" | "false") space',
    "number": '("-"? ([0-9] | [1-9] [0-9]*)) ("." [0-9]+)? ([eE] [-+]? [0-9]+)? space',
    "integer": '("-"? ([0-9] | [1-9] [0-9]*)) space',
    "string": r""" "\"" (
        [^"\\] |
        "\\" (["\\/bfnrt] | "u" [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F])
      )* "\"" space """,
    "null": '"null" space',
}

INVALID_RULE_CHARS_RE = re.compile(r"[^a-zA-Z0-9-]+")
GRAMMAR_LITERAL_ESCAPE_RE = re.compile(r'[\r\n"]')
GRAMMAR_LITERAL_ESCAPES = {"\r": "\\r", "\n": "\\n", '"': '\\"'}


class SchemaConverter:
    def __init__(self, prop_order):
        self._prop_order = prop_order
        self._rules = {"space": SPACE_RULE}

    def _format_literal(self, literal):
        escaped = GRAMMAR_LITERAL_ESCAPE_RE.sub(
            lambda m: GRAMMAR_LITERAL_ESCAPES.get(m.group(0)), json.dumps(literal)
        )
        return f'"{escaped}"'

    def _add_rule(self, name, rule):
        esc_name = INVALID_RULE_CHARS_RE.sub("-", name)
        if esc_name not in self._rules or self._rules[esc_name] == rule:
            key = esc_name
        else:
            i = 0
            while f"{esc_name}{i}" in self._rules:
                i += 1
            key = f"{esc_name}{i}"
        self._rules[key] = rule
        return key

    def visit(self, schema, name):
        schema_type = schema.get("type")
        rule_name = name or "root"

        if "oneOf" in schema or "anyOf" in schema:
            rule = " | ".join(
                (
                    self.visit(alt_schema, f'{name}{"-" if name else ""}{i}')
                    for i, alt_schema in enumerate(
                        schema.get("oneOf") or schema["anyOf"]
                    )
                )
            )
            return self._add_rule(rule_name, rule)

        elif "const" in schema:
            return self._add_rule(rule_name, self._format_literal(schema["const"]))

        elif "enum" in schema:
            rule = " | ".join((self._format_literal(v) for v in schema["enum"]))
            return self._add_rule(rule_name, rule)

        elif schema_type == "object" and "properties" in schema:
            # TODO: `required` keyword
            prop_order = self._prop_order
            prop_pairs = sorted(
                schema["properties"].items(),
                # sort by position in prop_order (if specified) then by key
                key=lambda kv: (prop_order.get(kv[0], len(prop_order)), kv[0]),
            )

            rule = '"{" space'
            for i, (prop_name, prop_schema) in enumerate(prop_pairs):
                prop_rule_name = self.visit(
                    prop_schema, f'{name}{"-" if name else ""}{prop_name}'
                )
                if i > 0:
                    rule += ' "," space'
                rule += rf' {self._format_literal(prop_name)} space ":" space {prop_rule_name}'
            rule += ' "}" space'

            return self._add_rule(rule_name, rule)

        elif schema_type == "array" and "items" in schema:
            # TODO `prefixItems` keyword
            item_rule_name = self.visit(
                schema["items"], f'{name}{"-" if name else ""}item'
            )
            rule = (
                f'"[" space ({item_rule_name} ("," space {item_rule_name})*)? "]" space'
            )
            return self._add_rule(rule_name, rule)

        else:
            assert schema_type in PRIMITIVE_RULES, f"Unrecognized schema: {schema}"
            return self._add_rule(
                "root" if rule_name == "root" else schema_type,
                PRIMITIVE_RULES[schema_type],
            )

    def format_grammar(self):
        return "\n".join((f"{name} ::= {rule}" for name, rule in self._rules.items()))


def json_schema_to_gbnf(schema: str, prop_order: Optional[List[str]] = None):
    prop_order = prop_order or []
    schema = json.loads(schema)
    prop_order = {name: idx for idx, name in enumerate(prop_order)}
    converter = SchemaConverter(prop_order)
    converter.visit(schema, "")
    return converter.format_grammar()