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 second(std.heap.page_allocator);
    const t = timer.lap() / 1000;

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

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

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

    var input = try parseInput(arena.allocator());

    var max: usize = 0;

    for (input) |a| {
        for (input) |b| {
            if (std.mem.eql(u8, a, b)) continue;

            var pos: usize = 1;
            const first = try parseSnail(arena.allocator(), &a, null, &pos);
            const ab = try first.add(arena.allocator(), &b);
            try reduce(arena.allocator(), ab);

            const m = magnitude(ab);
            if (m > max) {
                max = m;
            }
        }
    }
    return max;
}

fn parseInput(a: std.mem.Allocator) ![]Str {
    const input = @embedFile(path);
    var in = std.mem.tokenize(u8, input, "\n");

    var ret = std.ArrayList(Str).init(a);

    while (in.next()) |line| {
        try ret.append(line);
    }

    return ret.toOwnedSlice();
}

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 "day18b" {
    try std.testing.expectEqual(@as(usize, 4650), try second(std.testing.allocator));
}