const std = @import("std");

const PATH = "input/day18.txt";
const Str = []const u8;

pub fn first(allocator: ?std.mem.Allocator) anyerror!usize {
    _ = allocator;

    const file = @embedFile(PATH);
    var lines = std.mem.tokenize(u8, file, "\n");

    var sum: usize = 0;
    while (lines.next()) |line| {
        var pos: usize = 0;
        sum += parseLine(line, &pos);
    }

    return sum;
}

const Operator = enum {
    plus,
    prod,
};

fn parseLine(line: Str, pos: *usize) usize {
    var ret: usize = 0;
    var oper: Operator = undefined;

    while (pos.* < line.len) : (pos.* += 1) {
        switch (line[pos.*]) {
            ' ' => {},
            '*' => oper = .prod,
            '+' => oper = .plus,
            '0', '1', '2', '3', '4', '5', '6', '7', '8', '9' => {
                const val = std.fmt.parseUnsigned(usize, line[pos.* .. pos.* + 1], 10) catch unreachable;
                switch (oper) {
                    .plus => ret += val,
                    .prod => ret *= val,
                }
            },
            '(' => {
                pos.* += 1;
                switch (oper) {
                    .plus => ret += parseLine(line, pos),
                    .prod => ret *= parseLine(line, pos),
                }
            },
            ')' => break,
            else => unreachable,
        }
    }

    return ret;
}

pub fn second(allocator: ?std.mem.Allocator) anyerror!usize {
    var arena = std.heap.ArenaAllocator.init(allocator.?);
    defer arena.deinit();

    const file = @embedFile(PATH);
    var lines = std.mem.tokenize(u8, file, "\n");

    var sum: usize = 0;
    while (lines.next()) |line| {
        sum += try shuntingYard(arena.allocator(), line);
    }

    return sum;
}

fn lower(a: u8, b: u8) bool {
    const Associativity = enum {
        left,
        right,
    };
    const Oper = struct {
        prec: u1,
        assoc: Associativity,
    };

    const x: Oper = switch (a) {
        '*' => .{ .prec = 0, .assoc = .left },
        '+' => .{ .prec = 1, .assoc = .left },
        else => unreachable,
    };
    const y: Oper = switch (b) {
        '*' => .{ .prec = 0, .assoc = .left },
        '+' => .{ .prec = 1, .assoc = .left },
        else => unreachable,
    };

    if ((x.assoc == .left and x.prec <= y.prec) or
        (x.assoc == .right and x.prec < y.prec)) return true;

    return false;
}

fn shuntingYard(allocator: std.mem.Allocator, input: Str) !usize {
    var stack = std.ArrayList(u8).init(allocator);
    defer stack.deinit();

    // var output = std.ArrayList(u8).init(allocator);
    // defer output.deinit();

    var nums = std.ArrayList(usize).init(allocator);
    defer nums.deinit();

    for (input) |ch| {
        switch (ch) {
            ' ' => {},
            '*', '+' => {
                while (stack.popOrNull()) |op| {
                    if (op != '(' and lower(ch, op)) {
                        // try output.append(op);
                        switch (op) {
                            '*' => try nums.append(nums.pop() * nums.pop()),
                            '+' => try nums.append(nums.pop() + nums.pop()),
                            else => unreachable,
                        }
                    } else {
                        try stack.append(op);
                        break;
                    }
                }

                try stack.append(ch);
            },
            '(' => {
                try stack.append(ch);
            },
            ')' => {
                while (stack.popOrNull()) |op| {
                    if (op != '(') {
                        // try output.append(op);
                        switch (op) {
                            '*' => try nums.append(nums.pop() * nums.pop()),
                            '+' => try nums.append(nums.pop() + nums.pop()),
                            else => unreachable,
                        }
                    } else break;
                }
            },
            '0', '1', '2', '3', '4', '5', '6', '7', '8', '9' => {
                // try output.append(ch);
                try nums.append(try std.fmt.parseUnsigned(usize, &[_]u8{ch}, 10));
            },
            else => unreachable,
        }
    }

    while (stack.popOrNull()) |op| {
        // try output.append(op);
        switch (op) {
            '*' => try nums.append(nums.pop() * nums.pop()),
            '+' => try nums.append(nums.pop() + nums.pop()),
            else => unreachable,
        }
    }

    // for (output.items) |item| {
    //     std.debug.print("{c}", .{item});
    // }
    // std.debug.print("\n", .{});

    return nums.items[0];
}

test "day18a" {
    try std.testing.expectEqual(@as(usize, 29839238838303), try first(std.testing.allocator));
}

test "day18b" {
    try std.testing.expectEqual(@as(usize, 201376568795521), try second(std.testing.allocator));
}

test "Part1: no parentheses" {
    var pos: usize = 0;
    try std.testing.expectEqual(@as(usize, 71), parseLine("1 + 2 * 3 + 4 * 5 + 6", &pos));
}

test "Part1: parentheses" {
    var pos: usize = 0;
    try std.testing.expectEqual(@as(usize, 51), parseLine("1 + (2 * 3) + (4 * (5 + 6))", &pos));
}

test "Part2: no parens" {
    try std.testing.expectEqual(@as(usize, 231), try shuntingYard(std.testing.allocator, "1 + 2 * 3 + 4 * 5 + 6"));
}

test "Part2: parens" {
    try std.testing.expectEqual(@as(usize, 46), try shuntingYard(std.testing.allocator, "2 * 3 + (4 * 5)"));
}

test "Part2: multi-parens" {
    try std.testing.expectEqual(@as(usize, 23340), try shuntingYard(std.testing.allocator, "((2 + 4 * 9) * (6 + 9 * 8 + 6) + 6) + 2 + 4 * 2"));
}