#!/usr/bin/env python3
import sys
import time
from pathlib import Path

GENERATE = [
    (
        "Expr",
        [
            "Binary     -> lhs: ExprP, op: Token, rhs: ExprP",
            "Unary      -> op: Token, expr: ExprP",
            "Call       -> expr: ExprP, args: Vec<ExprP>, rparen: Token",
            "Get        -> expr: ExprP, name: Token",
            "Primary    -> token: Token",
            "Function   -> lparen: Token, params: Vec<(Token <COMMA> Option<ExprP>)>, return_type: Option<ExprP>, body: Vec<StmtP>, rbrace: Token",
        ],
    ),
    (
        "Stmt",
        [
            "Expr       -> expr: ExprP",
            "Assign     -> lhs: Token, rhs: ExprP",
            "Set        -> expr: ExprP, name: Token, rhs: ExprP",
            "Block      -> lbrace: Token, stmts: Vec<StmtP>, rbrace: Token",
            "Return     -> return_kw: Token, expr: Option<ExprP>",
            "If         -> if_kw: Token, condition: ExprP, then_branch: BlockStmt, else_branch: Vec<StmtP>",
        ],
    ),
]


class FileWriter:
    def __init__(self, path: Path) -> None:
        self.path = path
        self.fp = open(path, "w")

    def write(self, text: str = "") -> None:
        self.fp.write(text)

    def line(self, line: str = "") -> None:
        self.fp.write(line)
        self.fp.write("\n")


def define_visitor(wr: FileWriter, base: str, type_lines: list[str]):
    types = []
    for line in type_lines:
        parts = line.split("->")
        members = [p.strip().replace("<COMMA>", ",") for p in parts[1].split(",")]
        types += [(parts[0].strip(), members)]

    wr.line(f"pub trait {base}Visitor " + "{")

    for type, members in types:
        wr.line(
            f"    fn visit_{type.lower()}_{base.lower()}(&mut self, {base.lower()}: &{type}{base});"
        )
    wr.line("}")
    wr.line()


def define_ast(wr: FileWriter, base: str, type_lines: list[str]):
    types = []
    for line in type_lines:
        parts = line.split("->")
        members = [p.strip().replace("<COMMA>", ",") for p in parts[1].split(",")]
        types += [(parts[0].strip(), members)]

    visitor = f"{base}Visitor"

    wr.line(f"pub trait {base}: Debug + Any " + "{")
    wr.line(f"    fn accept(&self, visitor: &mut dyn {visitor});")
    wr.line("    fn as_any(self: Box<Self>) -> Box<dyn Any>;")
    wr.line("    fn as_any_ref(&self) -> &dyn Any;")
    wr.line("}")
    wr.line()
    wr.line(f"pub type {base}P = Box<dyn {base} + 'static>;")
    wr.line()

    for type, members in types:
        wr.line("#[derive(Debug)]")
        wr.line(f"pub struct {type}{base} " + "{")
        for member in members:
            wr.line(f"    pub {member},")
        wr.line("}")
        wr.line()
        wr.line(f"impl {base} for {type}{base} " + "{")
        wr.line(f"    fn accept(&self, visitor: &mut dyn {visitor})" + "{")
        wr.line(f"        visitor.visit_{type.lower()}_{base.lower()}(self);")
        wr.line("    }")
        wr.line()
        wr.line("    fn as_any(self: Box<Self>) -> Box<dyn Any> {")
        wr.line("        self")
        wr.line("    }")
        wr.line()
        wr.line("    fn as_any_ref(&self) -> &dyn Any {")
        wr.line("        self")
        wr.line("    }")
        wr.line("}")
        wr.line()


def main():
    if len(sys.argv) == 1:
        path = Path(__file__).parent / "../src/ast.rs"
    else:
        path = Path(sys.argv[1])

    path = path.resolve()
    print(f"Using {path}")

    wr = FileWriter(path)

    wr.line(
        "// This is an auto-generated file. Any changes made to this file may be overwritten."
    )
    wr.line("// This file was created at: " + time.strftime("%Y-%m-%d %H:%M:%S"))

    # Imports
    wr.line("use std::fmt::Debug;")
    wr.line("use std::any::Any;")
    wr.line()
    wr.line("use crate::token::Token;")
    wr.line()

    # Visitors
    for base, types in GENERATE:
        define_visitor(wr, base, types)

    # AST
    for base, types in GENERATE:
        define_ast(wr, base, types)


if __name__ == "__main__":
    main()