import textwrap
import os
import importlib
from lark import Lark

from agast import *
from aggen import *

p = Lark(open("grammar.lark").read(), start="program", parser="lalr")

if __name__ == "__main__":
    with open("test/arith.ag") as f:
        data = f.read()

    cst = p.parse(data)

    trans = Parser()
    ast = trans.transform(cst)

    res = gen(ast)
    print("Grammar:")
    print(res.parser_data)

    if not os.path.exists("gen"):
        os.makedirs("gen")
    with open("gen/arith.py", "w") as f:
        fmt_str = textwrap.dedent(
            """
            # This document is generated by agtest.

            __all__ = ["parse"]
            from typing import Generic, TypeVar, Optional, Callable, Dict, Any
            from lark import Lark, Transformer
            T = TypeVar('T')
            builtins: Dict[str, Any] = {{
                "parseInt": lambda s: int(s)
            }}

            class Thunk(Generic[T]):
                ''' A thunk represents a value that may be computed lazily. '''
                value: Optional[T]
                def __init__(self, func: Callable[[], T]):
                    self.func = func
                    self.value = None
                def get(self) -> T:
                    if self.value is None:
                        self.value = self.func()
                    return self.value

            parser = Lark('''
            {pd}
            ''', parser='lalr', start={starts}, debug=True)

            {ex}

            class Trans(Transformer[None]): {transdef}

            def parse(input: str, start: Optional[str] = None) -> Any:
                tree = parser.parse(input, start)
                trans = Trans()
                res = trans.transform(tree)
                return res
        """
        )
        f.write(
            fmt_str.format(
                pd=res.parser_data,
                ex=res.extra,
                starts=list(res.starts),
                transdef=res.transdef,
            )
        )

    mod = importlib.import_module("gen.arith")
    print(mod.parse("1 + 2 * 3", start="expr"))  # type: ignore