#!/usr/bin/env python

from collections import deque
from typing import List, Tuple, Union, Iterator, Optional

class IncrementalJSONParser:
    """
    Incremental JSON Parser for handling JSON data generated by LLMs.

    This parser returns partial string data for each token passed into it, if that token contains
    string value data. It yields tuples containing three fields: 1) A list describing the current
    path in the JSON object, including array indices, 2) A string describing the data type of the
    value (int, float, logical, stringpart, or string), and 3) the value data.

    Example JSON:
    {
        "price": 19.99,
        "itemNo": "3735272",
        "modulars" : [
            {
                "zone": "A",
                "section": 29,
                "position": 10
            },
            {
                "zone": "A",
                "section": 29,
                "position": 15
            }
        ]
    }

    Yielded tuples example (stringparts may vary depending on how chunks are divided):
    (['price'], 'float', 19.99)
    (['itemNo'], 'stringpart', '37')
    (['itemNo'], 'stringpart', '3527')
    (['itemNo'], 'stringpart', '2')
    (['itemNo'], 'string', '3735272')
    (['modulars', 0, 'zone'], 'stringpart', 'A')
    (['modulars', 0, 'zone'], 'string', 'A')
    (['modulars', 0, 'section'], 'int', 29)
    (['modulars', 0, 'position'], 'int', 10)
    (['modulars', 1, 'zone'], 'stringpart', 'A')
    (['modulars', 1, 'zone'], 'string', 'A')
    (['modulars', 1, 'section'], 'int', 29)
    (['modulars', 1, 'position'], 'int', 15)
    """

    def __init__(self, chunks: List[str] = []) -> None:
        """
        Initialize the IncrementalJSONParser.

        Parameters:
        chunks (list): Initial list of JSON string chunks to be parsed.
        """
        self.chunks: Iterator[str] = iter(chunks)
        self.sent_chunks: deque[str] = deque()
        self.result_tuples: deque[Tuple[List[Union[str, int]], str, Union[int, float, bool, str]]] = deque()
        self.closed: bool = False

        self.chunk: Optional[str] = None
        self.char_index: int = 0

        # Stores a flat representation of the object hierarchy at the location
        # being parsed.
        # NOTE: the parser uses this as a state variable to track whether it is
        #       in an object or an array.
        #
        # Example 1
        # JSON snippet: {"JSONExample" : ["this", "is", "an", "example"]}
        # Stack representation: "is" == ["JSONExample", 1]
        #
        # Example 2
        # JSON snippet: {"multiple" : {"hierarchial" : ["levels", "here"]}}
        # Stack representation: "levels" == ["multiple", "hierarchial", 0]
        self.stack: List[Union[str, int]] = []

        # Stores the current parsing context.
        self.context: str = 'awaiting_value'

        # This buffer is used for values so that they can be returned in their entirety.
        # It also stores string data for keys while they are being read in.
        self.buffer: str = ''

        # Buffer used to collect parts of string values. This allows the parser to yield partial
        # string data ('stringpart') for each segment, which is useful for processing incoming
        # JSON data chunks in an incremental manner, as strings may be split across chunks.
        self.stringpart_buffer: str = ''

        self.escape_sequences: dict[str, str] = {
            '"': '"',
            '\\': '\\',
            '/': '/',
            'b': '\b',
            'f': '\f',
            'n': '\n',
            'r': '\r',
            't': '\t'
        }

    def _get_next_chunk(self) -> None:
        """Retrieve the next chunk of JSON data to be processed."""
        self.chunk = next(self.chunks, None)
        if not self.chunk and self.sent_chunks:
            self.chunk = self.sent_chunks.popleft()

        self.char_index = 0

    def _process_next_char(self) -> None:
        """Process the next character in the current chunk based on the current context."""
        def handle_comma_or_ending_brace(char: str, stack: List[Union[str, int]]) -> str:
            if char == ',':
                if isinstance(stack[-1], str) or stack[-1] is None:
                    stack[-1] = None  # Reset the key placeholder
                    return 'awaiting_key_string'
                elif isinstance(stack[-1], int):
                    stack[-1] += 1
                    return 'awaiting_value'
            elif char == '}':
                if stack and (isinstance(stack[-1], str) or stack[-1] is None):
                    stack.pop()
                else:
                    raise ValueError("Bracket mismatch: Unexpected '}' encountered")
                if not stack:
                    return 'end'
                else:
                    return 'awaiting_comma_or_ending_brace'
            elif char == ']':
                if stack and isinstance(stack[-1], int):
                    stack.pop()
                else:
                    raise ValueError("Bracket mismatch: Unexpected ']' encountered")
                if not stack:
                    return 'end'
                else:
                    return 'awaiting_comma_or_ending_brace'
            elif char in ' \t\n\r':
                return 'awaiting_comma_or_ending_brace'
            else:
                raise ValueError(f"Unexpected character {char} while awaiting comma or ending brace")

        if self.chunk is None:
            return

        char = self.chunk[self.char_index]
        self.char_index += 1

        # Implement the logic based on the context
        if self.context == 'awaiting_key_string':
            if char in ' \t\n\r':
                return
            elif char == '"':
                self.context = 'key_string'
                self.buffer = ''
            elif char == '}':
                if self.stack:
                    self.stack.pop()
                self.context = 'awaiting_comma_or_ending_brace'
            else:
                raise ValueError(f"Unexpected character {char} while awaiting key")

        elif self.context == 'key_string':
            if char == '\\':
                self.context = 'key_escape'
            elif char == '"':
                if not self.buffer:
                    raise ValueError("Empty key string in data")
                if self.stack[-1] is None:
                    self.stack[-1] = self.buffer  # Replace None with the key
                else:
                    raise ValueError("Unexpected state: key placeholder not found")
                self.context = 'awaiting_key_value_separator'
            else:
                self.buffer += char

        elif self.context == 'key_escape':
            if char in self.escape_sequences:
                self.buffer += self.escape_sequences[char]
                self.context = 'key_string'
            elif char == 'u':
                self.unicode_buffer = ''
                self.context = 'key_unicode'
            else:
                raise ValueError(f"Unexpected escape character {char} in key string")

        elif self.context == 'awaiting_key_value_separator':
            if char in ' \t\n\r':
                return
            elif char == ':':
                self.context = 'awaiting_value'
            else:
                raise ValueError(f"Unexpected character {char} while awaiting key-value separator")

        elif self.context == 'awaiting_value':
            if char in ' \t\n\r':
                return
            elif char == '"':
                self.context = 'value_stringpart'
                self.stringpart_buffer = ''
                self.buffer = ''
            elif char in '0123456789-tfn':  # numerical or logical
                self.context = 'value_nonstring'
                self.buffer = char
            elif char == '{':
                self.stack.append(None)  # Push None as a placeholder for the key
                self.context = 'awaiting_key_string'
            elif char == '[':
                self.stack.append(0)
            elif char == ']' and isinstance(self.stack[-1], int):
                self.stack.pop()
                self.context = 'awaiting_comma_or_ending_brace'
            else:
                raise ValueError(f"Unexpected character {char} while awaiting value")

        elif self.context == 'value_stringpart':
            if char == '\\':
                self.context = 'value_escape'
            elif char == '"':
                if self.stringpart_buffer:
                    self.result_tuples.append((self.stack[:], 'stringpart', self.stringpart_buffer))
                self.result_tuples.append((self.stack[:], 'string', self.buffer))
                
                if self.stack and isinstance(self.stack[-1], str):
                    self.stack[-1] = None  # Reset the key placeholder
                self.context = 'awaiting_comma_or_ending_brace'
            else:
                self.stringpart_buffer += char
                self.buffer += char

        elif self.context == 'value_escape':
            if char in self.escape_sequences:
                self.stringpart_buffer += self.escape_sequences[char]
                self.buffer += self.escape_sequences[char]
                self.context = 'value_stringpart'
            elif char == 'u':
                self.unicode_buffer = ''
                self.context = 'value_unicode'
            else:
                raise ValueError(f"Unexpected escape character {char} in value string")

        elif self.context == 'value_nonstring':
            if char in ' \t\n\r,]}':
                if self.buffer in ('true', 'false', 'null'):
                    if self.buffer == 'true':
                        value = True
                    elif self.buffer == 'false':
                        value = False
                    else:
                        value = None
                    self.result_tuples.append((self.stack[:], 'logical', value))
                else:
                    if '.' in self.buffer:
                        # FIXME: non-finite floating-point numbers (especially
                        #        Python's keywords for them) are not supported
                        #        by the spec.
                        value = float(self.buffer)
                        self.result_tuples.append((self.stack[:], 'float', value))
                    else:
                        value = int(self.buffer)
                        self.result_tuples.append((self.stack[:], 'int', value))

                if self.stack and isinstance(self.stack[-1], str):
                    self.stack[-1] = None  # Reset the key placeholder
                self.context = handle_comma_or_ending_brace(char, self.stack)
            else:
                self.buffer += char

        elif self.context == 'key_unicode':
            if len(self.unicode_buffer) < 4 and char in '0123456789abcdefABCDEF':
                self.unicode_buffer += char
                if len(self.unicode_buffer) == 4:
                    self.buffer += chr(int(self.unicode_buffer, 16))
                    self.context = 'key_string'
            else:
                raise ValueError(f"Invalid unicode escape sequence in key string: \\u{self.unicode_buffer}{char}")

        elif self.context == 'value_unicode':
            if len(self.unicode_buffer) < 4 and char in '0123456789abcdefABCDEF':
                self.unicode_buffer += char
                if len(self.unicode_buffer) == 4:
                    unicode_char = chr(int(self.unicode_buffer, 16))
                    self.stringpart_buffer += unicode_char
                    self.buffer += unicode_char
                    self.context = 'value_stringpart'
            else:
                raise ValueError(f"Invalid unicode escape sequence in value string: \\u{self.unicode_buffer}{char}")

        elif self.context == 'awaiting_comma_or_ending_brace':
            self.context = handle_comma_or_ending_brace(char, self.stack)

        elif self.context == 'end':
            if char not in ' \t\n\r':
                raise ValueError(f"Unexpected character {char} after end of JSON")

    def __iter__(self) -> 'IncrementalJSONParser':
        return self

    def __next__(self) -> Tuple[List[Union[str, int]], str, Union[int, float, bool, str]]:
        """
        Return the next result tuple from the parser.

        Returns:
        tuple: A tuple containing the current path, data type, and value data.

        Raises:
        RuntimeError: If the parser is closed.
        StopIteration: If there are no more result tuples to yield.
        """
        if self.closed:
            raise RuntimeError("Cannot use parser after it has been closed")

        while not self.result_tuples:
            if not self.chunk or self.char_index >= len(self.chunk):
                if self.stringpart_buffer and self.context == 'value_stringpart':
                    self.result_tuples.append((self.stack[:], 'stringpart', self.stringpart_buffer))
                self.stringpart_buffer = ''
                self._get_next_chunk()

            if not self.chunk:
                raise StopIteration

            self._process_next_char()

        return self.result_tuples.popleft()

    def send(self, chunk: str) -> None:
        """
        Send a new chunk of JSON data to the parser.

        Parameters:
        chunk (str): A string containing a chunk of JSON data.
        """
        self.sent_chunks.append(chunk)

    def close(self) -> None:
        """
        Close the parser, checking for completeness and invalidating its state.

        Raises:
        RuntimeError: If the parser is already closed.
        ValueError: If the JSON data is incomplete.
        """
        if self.closed:
            raise RuntimeError("Parser already closed")

        if self.stack or self.context not in ['awaiting_value', 'awaiting_comma_or_ending_brace', 'end']:
            raise ValueError("Incomplete JSON data")

        # Invalidate the parser's state
        self.chunks = iter([])
        self.sent_chunks.clear()
        self.result_tuples.clear()
        self.chunk = None
        self.buffer = ''
        self.stringpart_buffer = ''
        self.stack = []

        # Mark the parser as closed
        self.closed = True

    # Context manager protocol

    def __enter__(self) -> 'IncrementalJSONParser':
        return self

    def __exit__(self, exc_type, exc_value, exc_tb) -> bool:
        self.close()
        return False


if __name__ == '__main__':
    import argparse
    import sys
    from typing import IO

    def parse_json_from_stream(stream: IO[str], chunk_size: int) -> None:
        parser = IncrementalJSONParser()
        while True:
            chunk = stream.read(chunk_size)
            if not chunk:
                break
            parser.send(chunk)
            for result in parser:
                print(result)
        parser.close()

    def main() -> None:
        parser = argparse.ArgumentParser(description="Incremental JSON Parser CLI")
        parser.add_argument('file', nargs='?', type=argparse.FileType('r'), default=sys.stdin,
                            help="JSON file to be parsed. If not provided, reads from stdin.")
        parser.add_argument('--chunk-size', type=int, default=1024,
                            help="Size of chunks to split the JSON data into (in bytes).")

        args = parser.parse_args()

        with args.file as f:
            parse_json_from_stream(f, args.chunk_size)

    main()
