const std = @import("std");
const tokens = @import("tokens.zig");
const Token = tokens.Token;

pub const NodeList = std.ArrayList(*Node);
pub const StmtList = std.ArrayList(*Stmt);
pub const ExprList = std.ArrayList(*Expr);
pub const ParamList = std.ArrayList(ParamDecl);
pub const ConstList = std.ArrayList(SingleConst);

// TODO convert FnCall to something like PrefixOp / InfixOp / SuffixOp
pub const NodeType = enum {
    Root,
    FnDecl,
    ConstDecl,
    Block,
    Expr,
    Stmt,
};

pub const ParamDecl = struct {
    name: Token,
    typ: Token,
};

pub const FnDecl = struct {
    func_name: Token,
    params: ParamList,
    return_type: Token,
    body: StmtList,
};

pub const SingleConst = struct {
    name: Token,
    expr: *Expr,
};

pub const BinaryExpr = struct {
    left: *Expr,
    op: Token,
    right: *Expr,
};

pub const UnaryExpr = struct {
    op: Token,
    right: *Expr,
};

pub const LiteralExpr = union(enum) {
    Bool: bool,
    Integer: []const u8,
    Float: []const u8,
    String: []const u8,
};

pub const AssignExpr = struct {
    name: Token,
    value: *Expr,
};

pub const ExprType = enum {
    Assign,

    // vardecls as expressions is a hack
    VarDecl,

    Binary,
    Unary,
    Literal,
    Variable,
    Grouping,
};

pub const VarDecl = struct {
    assign: AssignExpr,
    mutable: bool = false,
};

pub const Expr = union(ExprType) {
    Assign: AssignExpr,
    VarDecl: VarDecl,

    Binary: BinaryExpr,
    Unary: UnaryExpr,
    Literal: LiteralExpr,

    Variable: Token,
    Grouping: *Expr,
};

pub const IfBranch = std.ArrayList(Stmt);

pub const IfStmt = struct {
    condition: *Expr,
    then_branch: IfBranch,
    else_branch: IfBranch,
};

pub const Stmt = union(enum) {
    Expr: *Expr,
    Println: *Expr,

    If: IfStmt,

    pub fn mkPrintln(allocator: *std.mem.Allocator, expr: *Expr) !*Stmt {
        var println = try allocator.create(Stmt);
        println.* = Stmt{ .Println = expr };
        return println;
    }
};

pub const Node = union(NodeType) {
    Root: NodeList,
    FnDecl: FnDecl,
    ConstDecl: ConstList,

    Block: StmtList,

    Expr: *Expr,
    Stmt: *Stmt,
};

pub fn mkRoot(allocator: *std.mem.Allocator) !*Node {
    var node = try allocator.create(Node);
    node.* = Node{ .Root = NodeList.init(allocator) };
    return node;
}

fn printIdent(ident: usize) void {
    var i: usize = 0;
    while (i < ident) : (i += 1) {
        std.debug.warn("\t");
    }
}

fn print(ident: usize, comptime fmt: []const u8, args: ...) void {
    printIdent(ident);
    std.debug.warn(fmt, args);
}

pub fn printNode(node: *Node, ident: usize) void {
    switch (node.*) {
        .FnDecl => |decl| {
            print(ident, "(fn {} (", decl.func_name.lexeme);

            for (decl.params.toSlice()) |param| {
                std.debug.warn("({} {}) ", param.name.lexeme, param.typ.lexeme);
            }

            std.debug.warn(") (\n");

            for (decl.body.toSlice()) |stmt| {
                printIdent(ident + 1);
                printStmt(stmt);
                std.debug.warn("\n");
            }

            print(ident, ")\n");
        },

        .ConstDecl => |consts| {
            print(ident, "(const (\n");

            for (consts.toSlice()) |const_decl| {
                print(
                    ident + 1,
                    "({} ",
                    const_decl.name.lexeme,
                );

                printExpr(const_decl.expr);
                std.debug.warn(")\n");
            }

            print(ident, "))\n");
        },

        .Root => {
            for (node.Root.toSlice()) |child| {
                printNode(child, ident + 1);
            }
        },

        .Expr => |expr| {
            printIdent(ident);
            printExpr(expr);
            std.debug.warn("\n");
        },

        .Stmt => |stmt| blk: {
            printIdent(ident);
            printStmt(stmt);
            std.debug.warn("\n");
        },

        else => {
            print(ident, "unknown node: {}\n", node);
        },
    }
}

fn parenthetize(name: []const u8, exprs: []*Expr) void {
    std.debug.warn("({}", name);

    for (exprs) |expr| {
        std.debug.warn(" ");
        printExpr(expr);
    }

    std.debug.warn(")");
}

pub fn printExpr(expr: *Expr) void {
    switch (expr.*) {
        .Binary => |binary| parenthetize(binary.op.lexeme, &[_]*Expr{ binary.left, binary.right }),
        .Unary => |unary| parenthetize(unary.op.lexeme, &[_]*Expr{unary.right}),
        .Grouping => |expr_ptr| parenthetize("group", &[_]*Expr{expr_ptr}),

        .Literal => |literal| {
            switch (literal) {
                .Bool => |val| std.debug.warn("{}", val),
                .Integer => |val| std.debug.warn("{}", val),
                .Float => |val| std.debug.warn("{}", val),
                .String => |val| std.debug.warn("'{}'", val),
            }
        },

        .Variable => |token| std.debug.warn("{}", token.lexeme),

        .VarDecl => |decl| {
            if (decl.mutable) {
                std.debug.warn("(mut ");
            } else {
                std.debug.warn("(");
            }

            std.debug.warn("let {} ", decl.assign.name.lexeme);
            printExpr(decl.assign.value);
            std.debug.warn(")");
        },

        .Assign => |assign| {
            std.debug.warn("(set ");
            std.debug.warn("{} ", assign.name.lexeme);
            printExpr(assign.value);
            std.debug.warn(")");
        },

        else => std.debug.warn("UnknownExpr-{}", @tagName(expr.*)),
    }
}

pub fn printStmt(stmt: *Stmt) void {
    switch (stmt.*) {
        .Println => |expr| parenthetize("println", &[_]*Expr{expr}),
        .Expr => |expr| printExpr(expr),
        else => std.debug.warn("UnknownStmt-{}", @tagName(stmt.*)),
    }
}