from os.path import dirname
import sys

sys.path.append(f"{dirname(__file__)}/../src")
from party import regex, Parser, just, FwDeclaration

from typing import TypeVar, Any

# Utilities
# whitespace = regex(r"\s*")
whitespace = just(" ").or_else(just("\n")).repeat()

T = TypeVar("T")


def lexeme(p: Parser[T]) -> Parser[T]:
    return p << whitespace


# Punctuation
lbrace = lexeme(just("{"))
rbrace = lexeme(just("}"))
lbrack = lexeme(just("["))
rbrack = lexeme(just("]"))
colon = lexeme(just(":"))
comma = lexeme(just(","))

# Primitives
true = lexeme(just("true")).value(True)
false = lexeme(just("false")).value(False)
null = lexeme(just("null")).value(None)
number = lexeme(regex(r"-?(0|[1-9][0-9]*)([.][0-9]+)?([eE][+-]?[0-9]+)?")).map(float)
just_part = regex(r'[^"\\]+')
just_esc = just("\\") >> (
    just("\\")
    | just("/")
    | just('"')
    | just("b").value("\b")
    | just("f").value("\f")
    | just("n").value("\n")
    | just("r").value("\r")
    | just("t").value("\t")
    | regex(r"u[0-9a-fA-F]{4}").map(lambda s: chr(int(s[1:], 16)))
)
quoted = lexeme(just('"') >> (just_part | just_esc).repeat().map(lambda l: "".join(l)) << just('"'))

# Data structures
json_value = FwDeclaration[Any]()
object_pair = (quoted << colon) & json_value.p()
json_object = lbrace >> object_pair.sep_by(comma).map(dict) << rbrace
array = lbrack >> json_value.p().sep_by(comma) << rbrack

# Everything
json_value.define(quoted | number | json_object | array | true | false | null)
json_doc = whitespace >> json_value.p()


def test():
    assert (
        json_doc.parse(
            r"""
    {
        "int": 1,
        "just": "hello",
        "a list": [1, 2, 3],
        "escapes": "\n \u24D2",
        "nested": {"x": "y"},
        "other": [true, false, null]
    }
"""
        )
        == {
            "int": 1,
            "just": "hello",
            "a list": [1, 2, 3],
            "escapes": "\n ⓒ",
            "nested": {"x": "y"},
            "other": [True, False, None],
        }
    )


if __name__ == "__main__":
    from sys import stdin
    test()
    # print(repr(json_doc.parse(stdin.read())))