from typing import *
import textwrap
import re
import copy

from agast import *

global i
i = 0

class GenResult:
    def __init__(self, pd: str = "", ex: str = ""):
        self.parser_data = pd
        self.extra = ex

def gen(program: List[Decl]) -> GenResult:
    res = GenResult()
    def gen(prefix: str = "", suffix: str = "") -> str:
        global i
        presan = re.sub("[^0-9a-zA-Z]+", "_", prefix)
        sufsan = re.sub("[^0-9a-zA-Z]+", "_", suffix)
        i += 1
        return f"{presan}{i}{sufsan}"
    def v(name: str) -> str:
        return f"__ag_{name}"

    # builtins
    builtins: Dict[str, str] = {
        "parseInt": "",
    }

    # collect a list of name -> iface declarations
    ifaces: Dict[str, Iface] = dict(
        map(lambda c: (c.name, cast(Iface, c)),
        filter(lambda c: isinstance(c, Iface),
            program)))
    
    # list of node -> iface mappings
    what_ifaces: Dict[str, Set[str]] = dict()
    what_fields: Dict[str, Dict[str, str]] = dict()
    for node in filter(lambda c: isinstance(c, Node), program):
        node = cast(Node, node)
        # all_fields = dict()
        what_ifaces[node.name] = set(node.ifaces)
        this_fields = dict()
        for iface in node.ifaces:
            fields = ifaces[iface].fields
            for field in fields:
                if field.name in this_fields:
                    raise Exception("duplicate field name")
                this_fields[field.name] = field.ty
        what_fields[node.name] = this_fields
    print("what_ifaces:", what_ifaces)
    print("what_fields:", what_fields)

    # a high-level dictionary of productions; this has sub-productions
    # that should be further expanded at a later step before converting
    # into lark code
    productions_hi: Dict[str, Union[str, List[str]]] = dict()

    # TODO: this should probably not be inlined here, but i'll move it
    # out once i get more info into the 'env'
    def collect_required_thunks(env: List[Tuple[str, NodeRef]], expr: Expr) -> Dict[str, str]:
        names = dict(env)
        print(f"collect_required_thunks({expr})", expr.__class__)
        if isinstance(expr, ExprDot):
            return collect_required_thunks(env, expr.left)
        elif isinstance(expr, ExprMul):
            a = collect_required_thunks(env, expr.left)
            b = collect_required_thunks(env, expr.right)
            a.update(b)
            return a
        elif isinstance(expr, ExprAdd):
            a = collect_required_thunks(env, expr.left)
            b = collect_required_thunks(env, expr.right)
            a.update(b)
            return a
        elif isinstance(expr, ExprCall):
            return collect_required_thunks(env, expr.func)
        elif isinstance(expr, ExprName):
            if expr.name not in names and expr.name not in builtins:
                raise Exception(f"unbound name '{expr.name}'")
            return dict()
        raise Exception(f"unhandled {expr.__class__}")

    for node in filter(lambda c: isinstance(c, Node), program):
        node = cast(Node, node)
        n_class_name = gen(node.name)
        class_decl = textwrap.dedent(f"""
            class {v(n_class_name)}: pass
        """)
        res.extra += class_decl

        print(node.name, node.ifaces)

        for variant in node.variants:
            v_class_name = gen(f"{n_class_name}_var")
            class_decl = textwrap.dedent(f"""
                class {v(v_class_name)}({v(n_class_name)}):
                    ''' '''
                    pass
            """)
            res.extra += class_decl

            prod_name = gen(node.name)
            print(prod_name)

            # create an environment for checking the equations based on
            # the production
            env: List[Tuple[str, NodeRef]] = list()
            for sym in variant.prod:
                if isinstance(sym, SymRename):
                    env.append((sym.name, sym.ty))
            print(env)
            
            # for each of the equations, find out what the equation is
            # trying to compute, and generate a thunk corresponding to
            # that value.
            for eq in variant.equations:
                eq_name = gen(f"eq_{node.name}")
                thunk_name = gen(f"thunk_{node.name}")

                print("RHS", eq.rhs, eq.rhs.id)
                collect_required_thunks(copy.deepcopy(env), eq.rhs)

                func_impl = textwrap.dedent(f"""
                    def {eq_name}() -> None:
                        ''' {repr(eq)} '''
                        pass
                    def {thunk_name}() -> Thunk[None]:
                        return Thunk({eq_name})
                """)
                print(f"```py\n{func_impl}\n```")
                res.extra += func_impl

        # this is a "type alias" that connects it to one of the generated
        # names above
        res.extra += f"{node.name} = {v(n_class_name)}"

    return res