const std = @import("std");

const path = "data/day18/input.txt";
const explode_limit = 4;

const Str = []const u8;
const SnailType = u6;

const SnailItem = union(enum) {
    num: SnailType,
    pair: *SnailNumber,
};

const SnailNumber = struct {
    left: ?SnailItem = null,
    right: ?SnailItem = null,
    parent: ?*SnailNumber = null,
    fn add(self: *@This(), a: std.mem.Allocator, input: *const Str) !*SnailNumber {
        var pos: usize = 1;
        const next = try parseSnail(a, input, null, &pos);

        const root = try a.create(SnailNumber);
        root.left = SnailItem{ .pair = self };
        root.right = SnailItem{ .pair = next };
        root.parent = null; // segfaults without this line :-(

        self.parent = root;
        next.parent = root;

        return root;
    }
};

pub fn main() !void {
    var timer = try std.time.Timer.start();
    const ret = try first(std.heap.page_allocator);
    const t = timer.lap() / 1000;

    try std.testing.expectEqual(@as(usize, 3675), ret);

    std.debug.print("Day 18a result: {d} \t\ttime: {d}us\n", .{ ret, t });
}

pub fn first(allocator: ?std.mem.Allocator) !usize {
    const input = @embedFile(path);
    var in = std.mem.tokenize(u8, input, "\n");

    var arena = std.heap.ArenaAllocator.init(allocator.?);
    defer arena.deinit();

    var ret: *SnailNumber = undefined;

    var idx: usize = 0;
    while (in.next()) |line| : (idx += 1) {
        if (idx == 0) {
            var pos: usize = 1;
            ret = try parseSnail(arena.allocator(), &line, null, &pos);
        } else {
            ret = try ret.add(arena.allocator(), &line);
            try reduce(arena.allocator(), ret);
        }
    }

    return magnitude(ret);
}

fn reduce(a: std.mem.Allocator, sn: *SnailNumber) anyerror!void {
    if (checkExplode(a, sn, 0)) |node| {
        explode(a, node);
        try reduce(a, sn);
    } else if (checkSplit(a, sn)) |node| {
        try split(a, node);
        try reduce(a, sn);
    }
}

fn explode(a: std.mem.Allocator, node: *SnailNumber) void {
    defer a.destroy(node); // free up abandoned node

    explodeLeft(node);
    explodeRight(node);

    const par = node.parent.?;
    if (par.left.? == .pair and par.left.?.pair == node) {
        par.left = SnailItem{ .num = 0 };
    } else {
        par.right = SnailItem{ .num = 0 };
    }
}

fn explodeLeft(node: *SnailNumber) void {
    const left = node.left.?.num;

    var prev = node;
    var it = node.parent;
    while (it) |par| : ({
        prev = par;
        it = par.parent;
    }) {
        if (par.left) |l| {
            // check if it is not the exploding node
            if (l == .pair and l.pair == prev) continue;
            break;
        }
    }

    target: while (it) |target| {
        // These must be _pointers_! Removing the & makes them local, so the changes
        // do not happen when they should.
        for ([2]*?SnailItem{ &target.right, &target.left }) |t| {
            if (t.*) |*child| {
                switch (child.*) {
                    SnailItem.pair => |c| {
                        if (c != prev) {
                            it = c;
                            continue :target;
                        }
                    },
                    SnailItem.num => |*c| {
                        c.* += left;
                        return;
                    },
                }
            }
        }
    }
}

fn explodeRight(node: *SnailNumber) void {
    const right = node.right.?.num;

    var prev = node;
    var it = node.parent;
    while (it) |par| : ({
        prev = par;
        it = par.parent;
    }) {
        if (par.right) |r| {
            // check if it is not the exploding node
            if (r == .pair and r.pair == prev) continue;
            break;
        }
    }

    target: while (it) |target| {
        // These must be _pointers_! Removing the & makes them local, so the changes
        // do not happen when they should.
        for ([2]*?SnailItem{ &target.left, &target.right }) |t| {
            if (t.*) |*child| {
                switch (child.*) {
                    SnailItem.pair => |c| {
                        if (c != prev) {
                            it = c;
                            continue :target;
                        }
                    },
                    SnailItem.num => |*c| {
                        c.* += right;
                        return;
                    },
                }
            }
        }
    }
}

fn checkExplode(a: std.mem.Allocator, node: *SnailNumber, depth: u3) ?*SnailNumber {
    if (depth == explode_limit) return node;

    if (node.left) |l| {
        if (l == .pair) {
            const ret = checkExplode(a, l.pair, depth + 1);
            if (ret != null) return ret;
        }
    }

    if (node.right) |r| {
        if (r == .pair) {
            const ret = checkExplode(a, r.pair, depth + 1);
            if (ret != null) return ret;
        }
    }

    return null;
}

fn parseSnail(a: std.mem.Allocator, in: *const Str, parent: ?*SnailNumber, pos: *usize) anyerror!*SnailNumber {
    // allocate and set up new node
    const node = try a.create(SnailNumber);
    node.parent = parent;

    var right: bool = false;

    while (pos.* < in.len) : (pos.* += 1) {
        var ch = in.*[pos.*];
        switch (ch) {
            ']' => {
                break;
            },
            '[' => {
                pos.* += 1;
                if (right) {
                    node.right = SnailItem{ .pair = try parseSnail(a, in, node, pos) };
                } else {
                    node.left = SnailItem{ .pair = try parseSnail(a, in, node, pos) };
                }
            },
            ',' => {
                right = true;
            },
            else => {
                if (right) {
                    node.right = SnailItem{ .num = try std.fmt.parseUnsigned(SnailType, in.*[pos.* .. pos.* + 1], 10) };
                } else {
                    node.left = SnailItem{ .num = try std.fmt.parseUnsigned(SnailType, in.*[pos.* .. pos.* + 1], 10) };
                }
            },
        }
    }

    return node;
}

fn split(a: std.mem.Allocator, node: *SnailNumber) !void {
    if (node.left) |l| {
        if (l == .num and l.num > 9) {
            const new = try a.create(SnailNumber);
            new.left = SnailItem{ .num = l.num / 2 };
            new.right = SnailItem{ .num = l.num - l.num / 2 };
            new.parent = node;

            node.left = SnailItem{ .pair = new };
            return;
        }
    }
    if (node.right) |r| {
        if (r == .num and r.num > 9) {
            const new = try a.create(SnailNumber);
            new.left = SnailItem{ .num = r.num / 2 };
            new.right = SnailItem{ .num = r.num - r.num / 2 };
            new.parent = node;

            node.right = SnailItem{ .pair = new };
            return;
        }
    }
}

fn checkSplit(a: std.mem.Allocator, node: *SnailNumber) ?*SnailNumber {
    if (node.left) |l| {
        switch (l) {
            SnailItem.pair => |n| {
                const ret = checkSplit(a, n);
                if (ret != null) return ret;
            },
            SnailItem.num => |val| if (val > 9) return node,
        }
    }

    if (node.right) |r| {
        switch (r) {
            SnailItem.pair => |n| {
                const ret = checkSplit(a, n);
                if (ret != null) return ret;
            },
            SnailItem.num => |val| if (val > 9) return node,
        }
    }

    return null;
}

fn printSnailNumber(node: *SnailNumber) void {
    std.debug.print("[", .{});
    if (node.left) |l| {
        switch (l) {
            SnailItem.num => |value| std.debug.print("{d}", .{value}),
            SnailItem.pair => |n| printSnailNumber(n),
        }
    }
    std.debug.print(",", .{});
    if (node.right) |r| {
        switch (r) {
            SnailItem.num => |value| std.debug.print("{d}", .{value}),
            SnailItem.pair => |n| printSnailNumber(n),
        }
    }
    std.debug.print("]", .{});
}

fn magnitude(node: *SnailNumber) usize {
    var left: ?usize = null;
    var right: ?usize = null;

    if (node.left) |l| {
        switch (l) {
            SnailItem.num => |value| left = value,
            SnailItem.pair => |n| left = magnitude(n),
        }
    }
    if (node.right) |r| {
        switch (r) {
            SnailItem.num => |value| right = value,
            SnailItem.pair => |n| right = magnitude(n),
        }
    }

    if (left) |l| {
        if (right) |r| {
            return 3 * l + 2 * r;
        } else {
            return l;
        }
    }
    return right.?;
}

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

// test "magnitude" {
//     var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
//     const allocator = arena.allocator();
//     defer arena.deinit();

//     const in = [_]Str{
//         "[[9,1],[1,9]]",
//         "[[1,2],[[3,4],5]]",
//         "[[[[8,7],[7,7]],[[8,6],[7,7]]],[[[0,7],[6,6]],[8,7]]]",
//     };

//     const out = [_]usize{
//         129,
//         143,
//         3488,
//     };

//     for (in) |str, idx| {
//         var pos: usize = 1;
//         const root = try parseSnail(allocator, &str[0..], null, &pos);
//         try std.testing.expectEqual(out[idx], magnitude(root));
//     }
// }

// test "reduce" {
//     const tc = [_][2]Str{
//         .{ "[[[[[4,3],4],4],[7,[[8,4],9]]],[1,1]]", "[[[[0,7],4],[[7,8],[6,0]]],[8,1]]" },
//         .{ "[[[[0,[4,5]],[0,0]],[[[4,5],[2,6]],[9,5]]],[7,[[[3,7],[4,3]],[[6,3],[8,8]]]]]", "[[[[4,0],[5,4]],[[7,7],[6,0]]],[[8,[7,7]],[[7,9],[5,0]]]]" },
//         .{ "[[[[[7,7],[7,7]],[[8,7],[8,7]]],[[[7,0],[7,7]],9]],[[[[4,2],2],6],[8,7]]]", "[[[[8,7],[7,7]],[[8,6],[7,7]]],[[[0,7],[6,6]],[8,7]]]" },
//         .{ "[1,1]", "[1,1]" },
//     };

//     var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
//     const allocator = arena.allocator();
//     defer arena.deinit();

//     for (tc) |t| {
//         std.debug.print("\n\ninput:\t{s}", .{t[0]});
//         std.debug.print("\nwant:\t{s}\n", .{t[1]});
//         var pos: usize = 1;
//         const root = try parseSnail(allocator, &t[0][0..], null, &pos);
//         try reduce(allocator, root);
//         std.debug.print("got:\t", .{});
//         printSnailNumber(root);
//     }
// }

// test "explode" {
//     const tc = [_][2]Str{
//         .{ "[[[[[9,8],1],2],3],4]", "[[[[0,9],2],3],4]" },
//         .{ "[7,[6,[5,[4,[3,2]]]]]", "[7,[6,[5,[7,0]]]]" },
//         .{ "[[6,[5,[4,[3,2]]]],1]", "[[6,[5,[7,0]]],3]" },
//         .{ "[[3,[2,[1,[7,3]]]],[6,[5,[4,[3,2]]]]]", "[[3,[2,[8,0]]],[9,[5,[4,[3,2]]]]]" },
//         .{ "[[3,[2,[8,0]]],[9,[5,[4,[3,2]]]]]", "[[3,[2,[8,0]]],[9,[5,[7,0]]]]" },
//         .{ "[[[[[4,3],4],4],[7,[[8,4],9]]],[1,1]]", "[[[[0,7],4],[7,[[8,4],9]]],[1,1]]" },
//         .{ "[[[[0,7],4],[7,[[8,4],9]]],[1,1]]", "[[[[0,7],4],[15,[0,13]]],[1,1]]" },
//         .{ "", "" },
//     };

//     var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
//     const allocator = arena.allocator();
//     defer arena.deinit();

//     for (tc) |t| {
//         std.debug.print("\nwant:\t{s} input: {s}\n", .{ t[1], t[0] });
//         var pos: usize = 1;
//         const root = try parseSnail(allocator, &t[0][0..], null, &pos);
//         explode(allocator, checkExplode(allocator, root, 0).?);
//         std.debug.print("got:\t", .{});
//         printSnailNumber(root);
//         try std.testing.expect(1 == 1);
//     }
// }

// test "split" {
//     const tc = [_][2]Str{
//         .{ "[[[[0,7],4],[7,[[8,4],9]]],[1,1]]", "[[[[0,7],4],[[7,8],[0,13]]],[1,1]]" },
//     };

//     var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
//     const allocator = arena.allocator();
//     defer arena.deinit();

//     for (tc) |t| {
//         std.debug.print("\nwant:\t{s} input: {s}\n", .{ t[1], t[0] });
//         var pos: usize = 1;
//         const root = try parseSnail(allocator, &t[0][0..], null, &pos);
//         explode(allocator, checkExplode(allocator, root, 0).?);
//         try split(allocator, checkSplit(allocator, root).?);
//         std.debug.print("got:\t", .{});
//         printSnailNumber(root);
//         try std.testing.expect(1 == 1);
//     }
// }