const std = @import("std");
const Allocator = std.mem.Allocator;
const AstPrinter = @This();
const Expr = @import("expr.zig").Expr;
const Token = @import("Token.zig");

allocator: Allocator,

pub fn init(allocator: Allocator) AstPrinter {
    return .{
        .allocator = allocator,
    };
}

pub fn print(self: *const AstPrinter, expr: *const Expr) Allocator.Error![]const u8 {
    return switch (expr.*) {
        .binary => |binary| visitBinaryExpr(self, binary),
        .grouping => |grouping| visitGroupingExpr(self, grouping),
        .literal => |literal| visitLiteralExpr(self, literal),
        .unary => |unary| visitUnaryExpr(self, unary),
    };
}

fn visitBinaryExpr(self: *const AstPrinter, expr: Expr.Binary) ![]const u8 {
    return self.parenthesize(expr.operator.lexeme, &.{ expr.left, expr.right });
}

fn visitGroupingExpr(self: *const AstPrinter, expr: Expr.Grouping) ![]const u8 {
    return self.parenthesize("group", &.{expr.expression});
}

fn visitLiteralExpr(self: *const AstPrinter, expr: Expr.Literal) ![]const u8 {
    return switch (expr.value) {
        .string => |string| self.allocator.dupe(u8, string),
        .number => |number| std.fmt.allocPrint(self.allocator, "{}", .{number}),
        .boolean => |boolean| std.fmt.allocPrint(self.allocator, "{}", .{boolean}),
        .nil => self.allocator.dupe(u8, "nil"),
    };
}

fn visitUnaryExpr(self: *const AstPrinter, expr: Expr.Unary) ![]const u8 {
    return self.parenthesize(expr.operator.lexeme, &.{expr.right});
}

fn parenthesize(self: *const AstPrinter, name: []const u8, exprs: []const *const Expr) ![]const u8 {
    var builder: std.ArrayList(u8) = .empty;

    try builder.append(self.allocator, '(');
    try builder.appendSlice(self.allocator, name);
    for (exprs) |expr| {
        try builder.append(self.allocator, ' ');
        const printed = try print(self, expr);
        defer self.allocator.free(printed);
        try builder.appendSlice(self.allocator, printed);
    }
    try builder.append(self.allocator, ')');

    return builder.toOwnedSlice(self.allocator);
}

test {
    const allocator = std.testing.allocator;

    const expression = Expr{
        .binary = .init(
            &.{ .unary = .init(
                .init(.minus, "-", null, 1),
                &.{ .literal = .init(.{ .number = 123 }) },
            ) },
            .init(.star, "*", null, 1),
            &.{ .grouping = .init(
                &.{ .literal = .init(.{ .number = 45.67 }) },
            ) },
        ),
    };

    const ast_printer = AstPrinter.init(allocator);
    const printed = try ast_printer.print(&expression);
    defer allocator.free(printed);
    std.debug.print("{s}\n", .{printed});
}