const std = @import("std");

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

pub fn first(allocator: ?std.mem.Allocator) anyerror!usize {
    var memory = std.AutoHashMap(u36, u36).init(allocator.?);
    defer memory.deinit();

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

    var mask: Str = undefined;

    while (lines.next()) |line| {
        switch (line[1]) {
            // mask
            'a' => {
                var parts = std.mem.tokenize(u8, line, "=");
                _ = parts.next();
                mask = parts.next().?[1..];
            },
            // mem
            'e' => {
                // parse memory data
                var parts = std.mem.tokenize(u8, line, "[]=");
                _ = parts.next(); // mem
                const address = try std.fmt.parseUnsigned(u36, parts.next().?, 0);
                _ = parts.next(); // emptry
                var value = try std.fmt.parseUnsigned(u36, parts.next().?[1..], 0);

                // enforce mask
                var idx: u6 = 36;
                for (mask) |item| {
                    switch (item) {
                        'X' => {},
                        '1' => value |= @as(u36, 1) << idx - 1,
                        '0' => value &= ~(@as(u36, 1) << idx - 1),
                        else => unreachable,
                    }
                    idx -= 1;
                }

                // add to memory table
                try memory.put(address, value);
            },
            else => unreachable,
        }
    }

    var sum: usize = 0;
    var it = memory.valueIterator();
    while (it.next()) |item| {
        sum += item.*;
    }

    return sum;
}

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

    var memory = std.AutoHashMap(u36, u36).init(arena.allocator());
    defer memory.deinit();

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

    var mask: Str = undefined;

    while (lines.next()) |line| {
        switch (line[1]) {
            // mask
            'a' => {
                var parts = std.mem.tokenize(u8, line, "=");
                _ = parts.next();
                mask = parts.next().?[1..];
            },
            // mem
            'e' => {
                var addresses = std.ArrayList(u36).init(arena.allocator());
                defer addresses.deinit();

                // parse memory data
                var parts = std.mem.tokenize(u8, line, "[]=");
                _ = parts.next(); // mem
                const adr = try std.fmt.parseUnsigned(u36, parts.next().?, 0);
                try addresses.append(adr);
                _ = parts.next(); // emptry
                const value = try std.fmt.parseUnsigned(u36, parts.next().?[1..], 0);

                // enforce mask
                var idx: u6 = 36;
                for (mask) |item| {
                    defer idx -= 1;
                    switch (item) {
                        'X' => {
                            // This line below is _very_ important. Using append() would
                            // arbitrary change array size, which makes .items pointers
                            // unpredictable, resulting strange errors.
                            try addresses.ensureTotalCapacity(addresses.capacity * 2);
                            for (addresses.items) |*address| {
                                const zero = address.* & ~(@as(u36, 1) << idx - 1);
                                addresses.appendAssumeCapacity(zero);
                                const one = address.* | (@as(u36, 1) << idx - 1);
                                address.* = one;
                            }
                        },
                        '1' => {
                            for (addresses.items) |*address| {
                                address.* |= @as(u36, 1) << idx - 1;
                            }
                        },
                        '0' => {},
                        else => unreachable,
                    }
                }

                // add to memory table
                for (addresses.items) |address| {
                    try memory.put(address, value);
                }
            },
            else => unreachable,
        }
    }

    var sum: usize = 0;
    var it = memory.valueIterator();
    while (it.next()) |item| {
        sum += item.*;
    }

    return sum;
}

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

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

const test_input =
    \\mask = XXXXXXXXXXXXXXXXXXXXXXXXXXXXX1XXXX0X
    \\mem[8] = 11
    \\mem[7] = 101
    \\mem[8] = 0
;

const test_input2 =
    \\mask = 000000000000000000000000000000X1001X
    \\mem[42] = 100
    \\mask = 00000000000000000000000000000000X0XX
    \\mem[26] = 1
;