const std = @import("std");

pub fn main() anyerror!void {
    var arena = std.heap.ArenaAllocator.init(std.heap.page_allocator);
    defer arena.deinit();
    var gpa = &arena.allocator;

    var sum : i64 = 0;
    var f = try std.fs.cwd().openFile("input", .{});
    var contents = try f.readToEndAlloc(gpa, std.math.maxInt(u32));
    defer gpa.free(contents);
    var it = std.mem.tokenize(contents, "\n");
    while (it.next()) |line| {
        var tree = try Tree.init(gpa);
        defer tree.deinit();
        try tree.parse_line(line);
        sum += tree.resolve();
    }

    std.log.info("Sum of all operations from input: {}", .{sum});

    sum = 0;
    it = std.mem.tokenize(contents, "\n");
    while (it.next()) |line| {
        var tree = try Tree.init(gpa);
        defer tree.deinit();
        try tree.parse_line_part2(line);
        sum += tree.resolve();
    }

    std.log.info("Part2: Sum of all operations from input: {}", .{sum});
}

const Operation = enum {
    addition,
    multiplication,
    unknown,
};

const TreeValueTag = enum {
    value,
    operation,
};

const TreeValue = union(TreeValueTag) {
    value: i64,
    operation: Operation,
};

/// Binary tree node
const TreeElement = struct {
    parent: ?*TreeElement = null,
    children: [2]?*TreeElement = [_]?*TreeElement {null, null},
    value: TreeValue,

    pub fn add_child(self: *TreeElement, c: *TreeElement) void {
        //std.log.warn("Adding child {*} to {*}", .{c, self});
        for (self.children) |ch, k| {
            //std.log.warn("{}: {*} ({})", .{k, ch, ch == null});
            if (ch == null) {
                //std.log.warn("Set child {} of {*} to {*}", .{k, self, c});
                self.children[k] = c;
                return;
            }
        }
        // We're not meant to call add_child 3+ times on a TreeElement
        unreachable;
    }

    pub fn get_root_node(self: *TreeElement) *TreeElement {
        var r = self;
        while(r.parent) |p| {
            r = p;
        }
        return r;
    }

    pub fn print(self: *TreeElement, depth: u64) anyerror!void {
        var _type = switch(@as(TreeValueTag, self.*.value)) {
            .value => "value",
            .operation => "operation"
        };
        var i : u64 = 0;
        var stdout = std.io.getStdOut().writer();
        while (i < depth) : (i += 1) {
            try stdout.print(" ", .{});
        }
        try stdout.print("{*} ({}) - {}\n", .{self, _type, self.value});
        for (self.children) |c, k| {
            //std.log.warn("Child {} of {*} is {*}", .{k, self, c});
            if (c != null) {
                try TreeElement.print(self.children[k].?, depth + 2);
                //try self.children[k].?.print(depth + 2);
            }
        }
    }

    pub fn resolve(self: *TreeElement) i64 {
        if (@as(TreeValueTag, self.value) == TreeValueTag.value) {
            return self.value.value;
        }
        else {
            std.debug.assert(self.value.operation != .unknown);
            var v : i64 = undefined;
            var l = self.children[0].?.resolve();
            var r = self.children[1].?.resolve();
            if (self.value.operation == .addition) {
                v = l + r;
                std.log.warn("{} + {} = {}", .{l, r, v});
            }
            else {
                v = l * r;
                std.log.warn("{} * {} = {}", .{l, r, v});
            }
            return v;
        }
    }

    pub fn swap_child(self: *TreeElement, old: *TreeElement, new: *TreeElement) void {
        for (self.children) |v, k| {
            if (v == old) {
                self.children[k] = new;
                return;
            }
        }
        // maybe we don't want to fail at swapping a child right now
        unreachable;
    }
    pub fn nearest_operation(self: *TreeElement) ?*TreeElement {
        var t : ?*TreeElement = self;
        var r : ?*TreeElement = null;
        while (t) |node| {
            if (@as(TreeValueTag, node.value) == TreeValueTag.operation) {
                r = node;
                break;
            }
            else {
                t = node.parent;
            }
        }
        return r;
    }

    pub fn nearest_unknown_parent(self: *TreeElement) ?*TreeElement {
        var t : *TreeElement = self;
        var r : ?*TreeElement = null;
        while(t.parent) |p| {
            if (@as(TreeValueTag, p.*.value) == TreeValueTag.value) {
                    t = p;
                    continue;
            }
            else {
                r = p;
                break;
            }
        }
        return r;
    }
};

const Tree = struct {
    allocator: *std.mem.Allocator,
    // we want constant pointers for the TreeElements, but
    // maybe we can use this object to simplify keeping track
    // of them
    children: std.ArrayList(*TreeElement),

    pub fn init(a: *std.mem.Allocator) !*Tree {
        var self = try a.create(Tree);
        errdefer a.destroy(self);
        self.* = Tree {
            .allocator = a,
            .children = std.ArrayList(*TreeElement).init(a),
        };
        return self;
    }

    pub fn deinit(self: *Tree) void {
        for (self.children.items) |c| {
            self.allocator.destroy(c);
        }
        self.children.deinit();
        self.allocator.destroy(self);
    }

    pub fn resolve(self: *Tree) i64 {
        std.debug.assert(self.children.items.len > 0);
        var root = self.children.items[0].get_root_node();
        return root.resolve();
    }

    pub fn create_number_node(self: *Tree, buf: []const u8, previous: ?*TreeElement) !*TreeElement {
        var t = try self.create_node();
        var last : *TreeElement = undefined;
        t.* = .{
            .value = TreeValue { .value = try std.fmt.parseInt(i64, buf, 10) },
        };
        if (previous) |l| {
            l.add_child(t);
            t.parent = l;
            last = l;
        }
        else {
            last = t;
        }
        return last;
    }

    pub fn parse_line(self: *Tree, line: []const u8) !void {
        var last_node : ?*TreeElement = null;
        var disjointed_nodes = std.ArrayList(*TreeElement).init(self.allocator);
        defer disjointed_nodes.deinit();

        var buf = std.mem.zeroes([16]u8);
        var buf_pos : usize = 0;
        for (line) |c, k| {
            if (c == ' ') {
                // Check buffer length, if not zero add the necessary tree node
                if (buf_pos == 0) {
                    continue;
                }
                last_node = try self.create_number_node(buf[0..buf_pos], last_node);
                // Reset buffer
                std.mem.set(u8, buf[0..], 0);
                buf_pos = 0;
                continue;
            }
            if (std.ascii.isDigit(c)) {
                buf[buf_pos] = c;
                buf_pos += 1;

                // If we're at the end of the string, we need to run the stuff
                // for adding a number node
                if (k == line.len-1) {
                    last_node = try self.create_number_node(buf[0..buf_pos], last_node);
                    // Buffer reset unneccesary here
                }
                continue;
            }
            switch (c) {
                '+' => {
                    std.debug.assert(last_node != null);
                    // If we have a last node, we should check up for
                    // an unknown parent. If an unknown parent exists,
                    // we should set the operation on that node, and
                    // set that node to the last_node.
                    // If we ourselves are unknown
                    if (@as(TreeValueTag, last_node.?.value) == .operation) {
                        if (last_node.?.value.operation == .unknown) {
                            last_node.?.value = TreeValue { .operation = .addition };
                            continue;
                        }
                    }
                    else if (last_node.?.nearest_unknown_parent()) |unknown_parent| {
                        unknown_parent.value = TreeValue { .operation = .addition };
                        last_node = unknown_parent;
                        continue;
                    }
                    // If there is no unknown parent, our last node could
                    // be an operation node or value node.
                    // When it's an operation node, this is probably an error?
                    var root = last_node.?.get_root_node();
                    var t = try self.create_node();
                    t.* = .{
                        .value = TreeValue{ .operation = .addition},
                    };
                    root.parent = t;
                    t.add_child(root);
                    last_node = t;
                    continue;
                },
                '*' => {
                    std.debug.assert(last_node != null);
                    // Same as '+'
                    if (@as(TreeValueTag, last_node.?.value) == .operation) {
                        if (last_node.?.value.operation == .unknown) {
                            last_node.?.value = TreeValue { .operation = .multiplication };
                            continue;
                        }
                    }
                    else if (last_node.?.nearest_unknown_parent()) |unknown_parent| {
                        unknown_parent.value = TreeValue { .operation = .multiplication };
                        last_node = unknown_parent;
                        continue;
                    }
                    var root = last_node.?.get_root_node();
                    var t = try self.create_node();
                    t.* = .{
                        .value = TreeValue{ .operation = .multiplication},
                    };
                    root.parent = t;
                    t.add_child(root);
                    last_node = t;
                    continue;
                },
                '(' => {
                    // We need to start a new disjointed tree.
                    // Create an operation node of an unknown type,
                    // and set that as the parent of the current last node, if
                    // there is one.
                    if (last_node) |l| {
                        try disjointed_nodes.append(l);
                    }
                    else {
                        var t = try self.create_node();
                        t.* = .{
                            .value = TreeValue { .operation = .unknown },
                        };
                        if (last_node) |l| {
                            l.parent = t;
                            t.add_child(l);
                        }
                        try disjointed_nodes.append(t);
                    }
                    last_node = null;
                },
                ')' => {
                    std.debug.assert(last_node != null);
                    // Close out any number that's in progress
                    if (buf_pos != 0) {
                        last_node = try self.create_number_node(buf[0..buf_pos], last_node);
                        // Reset buffer
                        std.mem.set(u8, buf[0..], 0);
                        buf_pos = 0;
                    }
                    var dj = disjointed_nodes.pop();
                    if (last_node) |l| {
                        l.parent = dj;
                        dj.add_child(l);
                        last_node = dj;
                    }
                },
                else => unreachable,
            }
        }
    }

    // Addition hash priority over multiplication
    pub fn parse_line_part2(self: *Tree, line: []const u8) !void {
        var last_node : ?*TreeElement = null;
        var disjointed_nodes = std.ArrayList(*TreeElement).init(self.allocator);
        defer disjointed_nodes.deinit();

        var buf = std.mem.zeroes([16]u8);
        var buf_pos : usize = 0;
        for (line) |c, k| {
            if (c == ' ') {
                // Check buffer length, if not zero add the necessary tree node
                if (buf_pos == 0) {
                    continue;
                }
                last_node = try self.create_number_node(buf[0..buf_pos], last_node);
                // Reset buffer
                std.mem.set(u8, buf[0..], 0);
                buf_pos = 0;
                continue;
            }
            if (std.ascii.isDigit(c)) {
                buf[buf_pos] = c;
                buf_pos += 1;

                // If we're at the end of the string, we need to run the stuff
                // for adding a number node
                if (k == line.len-1) {
                    last_node = try self.create_number_node(buf[0..buf_pos], last_node);
                    // Buffer reset unneccesary here
                }
                continue;
            }
            switch (c) {
                '+' => {
                    std.debug.assert(last_node != null);
                    // If we have a last node, we should check up for
                    // an unknown parent. If an unknown parent exists,
                    // we should set the operation on that node, and
                    // set that node to the last_node.
                    // If we ourselves are unknown
                    if (@as(TreeValueTag, last_node.?.value) == .operation) {
                        if (last_node.?.value.operation == .unknown) {
                            last_node.?.value = TreeValue { .operation = .addition };
                            continue;
                        }
                    }
                    else if (last_node.?.nearest_unknown_parent()) |unknown_parent| {
                        unknown_parent.value = TreeValue { .operation = .addition };
                        last_node = unknown_parent;
                        continue;
                    }
                    // If there is no unknown parent, our last node could
                    // be an operation node or value node.
                    // If last_node has a parent which is an multiplication node,
                    // we want to insert ourselves between, otherwise we go to the top
                    // like in part one.
                    // If the last-node is itself a multiplication node, then we
                    // want to do the swap as well.
                    var t = try self.create_node();
                    var nearest_op = last_node.?.nearest_operation();
                    if (nearest_op) |p| {
                        var new_child : *TreeElement = undefined;
                        if (p.children[1] != null) {
                            new_child = p.children[1].?;
                        }
                        else if (p.children[0] != null) {
                            new_child = p.children[0].?;
                        }
                        else {
                            // maybe?
                            unreachable;
                        }
                        p.swap_child(new_child, t);
                        new_child.parent = t;
                        t.add_child(new_child);
                        t.parent = p;
                        last_node = t;
                    }
                    else {
                        if (@as(TreeValueTag, last_node.?.value) == TreeValueTag.operation and
                                last_node.?.value.operation == .multiplication) {
                            std.debug.assert(last_node.?.children[1] != null);
                            var mult_child = last_node.?.children[1].?;
                            last_node.?.swap_child(mult_child, t);
                            mult_child.parent = t;
                            t.add_child(mult_child);
                            t.parent = last_node.?;
                            last_node = t;
                        }
                        else {
                            var root = last_node.?.get_root_node();
                            root.parent = t;
                            t.add_child(root);
                            last_node = t;
                        }
                    }
                    t.value = TreeValue{ .operation = .addition};
                    continue;
                },
                '*' => {
                    std.debug.assert(last_node != null);
                    // Same as '+'
                    if (@as(TreeValueTag, last_node.?.value) == .operation) {
                        if (last_node.?.value.operation == .unknown) {
                            last_node.?.value = TreeValue { .operation = .multiplication };
                            continue;
                        }
                    }
                    else if (last_node.?.nearest_unknown_parent()) |unknown_parent| {
                        unknown_parent.value = TreeValue { .operation = .multiplication };
                        last_node = unknown_parent;
                        continue;
                    }
                    var root = last_node.?.get_root_node();
                    var t = try self.create_node();
                    t.* = .{
                        .value = TreeValue{ .operation = .multiplication},
                    };
                    root.parent = t;
                    t.add_child(root);
                    last_node = t;
                    continue;
                },
                '(' => {
                    // We need to start a new disjointed tree.
                    // Create an operation node of an unknown type,
                    // and set that as the parent of the current last node, if
                    // there is one.
                    if (last_node) |l| {
                        try disjointed_nodes.append(l);
                    }
                    else {
                        // This assumption that the new disjointed node will
                        // be the new root is not always true since there is
                        // operator precedence now?
                        // We also sometimes loose nodes because of the swapping?
                        var t = try self.create_node();
                        t.value = TreeValue { .operation = .unknown };
                        if (last_node) |l| {
                            l.parent = t;
                            t.add_child(l);
                        }
                        try disjointed_nodes.append(t);
                    }
                    last_node = null;
                },
                ')' => {
                    std.debug.assert(last_node != null);
                    // Close out any number that's in progress
                    if (buf_pos != 0) {
                        last_node = try self.create_number_node(buf[0..buf_pos], last_node);
                        // Reset buffer
                        std.mem.set(u8, buf[0..], 0);
                        buf_pos = 0;
                    }
                    var dj = disjointed_nodes.pop();
                    if (last_node) |l| {
                        var root = l.get_root_node();
                        root.parent = dj;
                        dj.add_child(root);
                        last_node = dj.get_root_node();
                    }
                },
                else => unreachable,
            }
        }
    }

    fn create_node(self: *Tree) !*TreeElement {
        var t = try self.allocator.create(TreeElement);
        // Having some non-zero initialization problems, but
        // can't use std.mem.zeroes since TreeValue doesn't
        // have a "zero" value.
        t.children = [_]?*TreeElement {null, null};
        t.parent = null;
        try self.children.append(t);
        return t;
    }

    fn print_tree(self: *Tree) !void {
        var stdout = std.io.getStdOut().writer();
        if (self.children.items.len == 0) {
            try stdout.print("Tree {} has no children", .{self});
            return;
        }
        var root = self.children.items[0].get_root_node();
        try root.print(0);
    }
};

test "no_parentheses" {
    var line = "1 + 2 * 3 + 4 * 5 + 6";
    var tree = try Tree.init(std.testing.allocator);
    defer tree.deinit();

    try tree.parse_line(line);
    std.debug.warn("\n", .{});
    try tree.print_tree();

    var v = tree.resolve();
    std.testing.expectEqual(@as(i64, 71), v);
}

test "ex2" {
    var line = "1 + (2 * 3) + (4 * (5 + 6))";
    var tree = try Tree.init(std.testing.allocator);
    defer tree.deinit();

    try tree.parse_line(line);
    std.debug.warn("\n", .{});
    try tree.print_tree();

    var v = tree.resolve();
    std.testing.expectEqual(@as(i64, 51), v);
}

test "ex3" {
    var line = "2 * 3 + (4 * 5)";
    var tree = try Tree.init(std.testing.allocator);
    defer tree.deinit();

    try tree.parse_line(line);
    std.debug.warn("\n", .{});
    try tree.print_tree();

    var v = tree.resolve();
    std.testing.expectEqual(@as(i64, 26), v);
}

test "ex4" {
    var line = "5 + (8 * 3 + 9 + 3 * 4 * 3)";
    var tree = try Tree.init(std.testing.allocator);
    defer tree.deinit();

    try tree.parse_line(line);
    std.debug.warn("\n", .{});
    try tree.print_tree();

    var v = tree.resolve();
    std.testing.expectEqual(@as(i64, 437), v);
}
test "ex5" {
    var line = "5 * 9 * (7 * 3 * 3 + 9 * 3 + (8 + 6 * 4))";
    var tree = try Tree.init(std.testing.allocator);
    defer tree.deinit();

    try tree.parse_line(line);
    std.debug.warn("\n", .{});
    try tree.print_tree();

    var v = tree.resolve();
    std.testing.expectEqual(@as(i64, 12240), v);
}

test "double_open_brackets" {
    var line = "((2 * 3) + 4) * 2";
    var tree = try Tree.init(std.testing.allocator);
    defer tree.deinit();

    try tree.parse_line(line);
    std.debug.warn("\n", .{});
    try tree.print_tree();

    var v = tree.resolve();
    std.testing.expectEqual(@as(i64, 20), v);
}

test "ex6" {
    var line = "((2 + 4 * 9) * (6 + 9 * 8 + 6) + 6) + 2 + 4 * 2";
    var tree = try Tree.init(std.testing.allocator);
    defer tree.deinit();

    try tree.parse_line(line);
    std.debug.warn("\n", .{});
    try tree.print_tree();

    var v = tree.resolve();
    std.testing.expectEqual(@as(i64, 13632), v);
}

fn test_part2(line: []const u8) !i64 {
    var tree = try Tree.init(std.testing.allocator);
    defer tree.deinit();

    try tree.parse_line_part2(line);
    std.debug.warn("\n", .{});
    try tree.print_tree();

    return tree.resolve();
}

test "part2_no_parentheses" {
    var line = "1 + 2 * 3 + 4 * 5 + 6";
    var v = try test_part2(line);
    std.testing.expectEqual(@as(i64, 231), v);
}

test "part2_ex2" {
    var line = "1 + (2 * 3) + (4 * (5 + 6))";
    var v = try test_part2(line);
    std.testing.expectEqual(@as(i64, 51), v);
}

test "part2_ex3" {
    var line = "2 * 3 + (4 * 5)";
    var v = try test_part2(line);
    std.testing.expectEqual(@as(i64, 46), v);
}

test "part2_ex4" {
    var line = "5 + (8 * 3 + 9 + 3 * 4 * 3)";
    var v = try test_part2(line);
    std.testing.expectEqual(@as(i64, 1445), v);
}

test "part2_ex5_sub" {
    var line = "5 * (2 * 3 + (4 * 5))";
    var v= try test_part2(line);
    std.testing.expectEqual(@as(i64, 230), v);
}

test "part2_ex5" {
    var line = "5 * 9 * (7 * 3 * 3 + 9 * 3 + (8 + 6 * 4))";
    var v = try test_part2(line);
    std.testing.expectEqual(@as(i64, 669060), v);
}

test "part2_ex6_sub" {

}
test "part2_ex6" {
    var line = "((2 + 4 * 9) * (6 + 9 * 8 + 6) + 6) + 2 + 4 * 2";
    var v = try test_part2(line);
    std.testing.expectEqual(@as(i64, 23340), v);
}