#!/usr/bin/python3

import re
import pathlib

root_dir = pathlib.Path(__file__).parent.parent
template_path = root_dir / "src/ast/node-template"
node_header_path = root_dir / "src/ast/node.h"
node_source_path = root_dir / "src/ast/node.cpp"
visitor_header_path = root_dir / "src/ast/visitor.h"
visitor_include_path = root_dir / "src/ast/visitorincl"

structs = {}


class Struct:
    name = ""
    parent = ""
    abstract = False
    fields = []

    def __init__(self, name, parent, abstract, fields):
        self.name = name
        self.parent = parent
        self.abstract = abstract
        self.fields = fields


def get_sections(file_name):
    with open(file_name, "r") as file:
        for word in file.read().split(";")[:-1]:
            yield word.strip()


def parse_section(section):
    global structs
    match = re.findall(
        r"(abstract )?(\w+)(?: : (\w+))?(?:\s*\{\s*(.+)\})?", section, re.DOTALL
    )
    if match:
        abstract = True if match[0][0] else False
        name = match[0][1]
        parent = match[0][2] or None
        fields = [field.strip().split() for field in match[0][3].splitlines()]

        structs[name] = Struct(name, parent, abstract, fields)


def get_all_fields(struct):
    return (
        get_all_fields(structs[struct.parent]) + struct.fields
        if struct.parent in structs
        else struct.fields
    )


def write_node_header():
    print(f"Writing to '{node_header_path}'")
    with open(node_header_path, "w") as file:
        file.write("// generated by nodegen.py; do not edit\n")
        file.write("\n")
        file.write("#pragma once\n")
        file.write("\n")
        file.write('#include "visitor.h"\n')
        file.write('#include "syntax/token.h"\n')
        file.write('#include "semantic/symdata.h"\n')
        file.write("\n")
        file.write("#include <string_view>\n")
        file.write("#include <memory>\n")
        file.write("#include <vector>\n")
        file.write("\n")

        def write_struct(struct):
            name = struct.name
            parent = struct.parent
            abstract = struct.abstract
            fields = struct.fields
            all_fields = get_all_fields(struct)

            # derived
            if parent:
                file.write(f"struct {name} : public {parent}\n")
            # base
            else:
                file.write(f"struct {name}\n")
            file.write("{\n")

            # fields
            for field in map(" ".join, fields):
                file.write(f"\t{field};\n")

            # constructor
            if name != "Node":
                # automatic constructor
                if not fields:
                    # not necessary to add constructor when parents have no fields
                    if struct.parent in structs and get_all_fields(
                        structs[struct.parent]
                    ):
                        file.write(f"\tusing {parent}::{parent};\n")
                # manual constructor
                else:
                    joined_fields = map(" ".join, all_fields)
                    file.write(f"\t{name}({', '.join(joined_fields)});\n")

            # define 'accept' for Node
            if name == "Node":
                file.write("\tvirtual void accept(Visitor &v) = 0;\n")
            # override 'accept'
            elif not abstract:
                file.write("\tvoid accept(Visitor &v) override;\n")

            file.write("};\n")

        for struct in structs.values():
            write_struct(struct)


def write_node_source():
    print(f"Writing to '{node_source_path}'")
    with open(node_source_path, "w") as file:
        file.write("// generated by nodegen.py; do not edit\n")
        file.write("\n")
        file.write('#include "node.h"\n')
        file.write("\n")

        def init_field(field):
            type = field[0]
            name = field[1]
            if "unique_ptr" in type:
                return f"std::move({name})"
            else:
                return f"{name}"

        def write_struct(struct):
            name = struct.name
            abstract = struct.abstract
            fields = struct.fields
            all_fields = get_all_fields(struct)

            # implement constructor
            if fields:
                # inherited fields
                parent_fields = [p for p in all_fields if p not in fields]
                # fields unique to current struct
                unique_fields = [p for p in all_fields if p in fields]

                # signature
                joined_params = list(
                    map(" ".join, parent_fields + unique_fields)
                )
                file.write(f"{name}::{name}({', '.join(joined_params)})\n")

                # initializer list
                parent_arguments = list(map(init_field, parent_fields))
                unique_arguments = list(map(init_field, unique_fields))

                initializer_list = []
                initializer_list.extend(
                    [
                        f"{struct.parent}({arg})"
                        for arg in parent_arguments
                    ]
                )
                initializer_list.extend(
                    [
                        f"{field[1]}({arg})"
                        for (field, arg) in zip(unique_fields, unique_arguments)
                    ]
                )
                file.write(f"\t: {', '.join(initializer_list)}\n")

                # body
                file.write("{\n")
                file.write("}\n")

            # implement 'accept'
            if not abstract:
                file.write(f"void {name}::accept(Visitor &v)\n")
                file.write("{\n")
                file.write("\tv.visit(*this);\n")
                file.write("}\n")

        for struct in structs.values():
            write_struct(struct)


def write_visitor_header():
    print(f"Writing to '{visitor_header_path}'")
    with open(visitor_header_path, "w") as file:
        file.write("// generated by nodegen.py; do not edit\n")
        file.write("\n")
        file.write("#pragma once\n")
        file.write("\n")

        def write_declaration(struct):
            if not struct.abstract:
                file.write(f"struct {struct.name};\n")

        for struct in structs.values():
            write_declaration(struct)

        file.write("\n")
        file.write("class Visitor\n")
        file.write("{\n")
        file.write("public:\n")

        def write_function(struct):
            if not struct.abstract:
                file.write(f"\tvirtual void visit({struct.name} &node) = 0;\n")

        for struct in structs.values():
            write_function(struct)

        file.write("};\n")


def write_visitor_include_header():
    print(f"Writing to '{visitor_include_path}'")
    with open(visitor_include_path, "w") as file:
        file.write("// generated by nodegen.py; do not edit\n")
        file.write("\n")

        def write_struct(struct):
            if not struct.abstract:
                file.write(f"void visit({struct.name} &node) override;\n")

        for struct in structs.values():
            write_struct(struct)


sections = get_sections(template_path)
for section in sections:
    parse_section(section)


write_node_header()
write_node_source()
write_visitor_header()
write_visitor_include_header()
