const std = @import("std");

const Lens = struct { label: []const u8, focalLength: u8 };

pub fn main() !void {
    var content = @embedFile("data.txt");

    var words = std.mem.split(u8, content, ",");

    var gpa = std.heap.GeneralPurposeAllocator(.{}){};
    const allocator = gpa.allocator();

    var boxes = std.AutoHashMap(u32, std.ArrayList(Lens)).init(allocator);
    defer boxes.deinit();

    var part1: u32 = 0;
    while (words.next()) |w| {
        // stupid last char thing
        const word = if (w[w.len - 1] == 10) w[0 .. w.len - 2] else w;

        // Part 1
        part1 += hash(word);

        // Part 2
        const dash = std.mem.indexOf(u8, word, "-");
        const equals = std.mem.indexOf(u8, word, "=");
        const idx = dash orelse equals orelse unreachable;
        const label = word[0..idx];
        const box = hash(label);
        std.debug.print("\nAfter \"{s}\":\n", .{word});

        if (dash != null) {
            // if dash, find the box
            if (boxes.getPtr(box)) |list| {
                // get the list, and find a lens with this label in it
                var indexToRemove: ?usize = null;
                for (list.items, 0..) |l, i| {
                    if (std.mem.eql(u8, l.label, label)) {
                        indexToRemove = i;
                    }
                }

                // if it's there, remove it
                if (indexToRemove) |i| {
                    _ = list.orderedRemove(i);
                }
            }
        } else {
            // if equals, parse out parts, put it in the box
            const focalLength = try std.fmt.parseInt(u8, word[idx + 1 ..], 10);
            const lens = Lens{ .focalLength = focalLength, .label = label };

            // does this list exist? use it. otherwise make a new one
            var list = boxes.get(box) orelse std.ArrayList(Lens).init(allocator);

            // find if this lens is already in the list
            var dupIdx: ?usize = null;
            for (list.items, 0..) |l, i| {
                if (std.mem.eql(u8, l.label, label)) {
                    dupIdx = i;
                }
            }

            // if it is, replace it with the new lens
            if (dupIdx) |i| {
                try list.replaceRange(i, 1, &[_]Lens{lens});
            } else {
                try list.append(lens);
            }
            try boxes.put(box, list);
        }

        printBoxes(boxes);
    }

    std.debug.print("Part 1: {d}\n", .{part1});

    var part2: usize = 0;
    var it = boxes.iterator();
    while (it.next()) |box| {
        const boxNumber = box.key_ptr.*;

        for (box.value_ptr.items, 0..) |lens, slotNumber| {
            part2 += (boxNumber + 1) * (slotNumber + 1) * lens.focalLength;
        }
    }

    std.debug.print("Part 2: {d}\n", .{part2});
}

fn hash(word: []const u8) u32 {
    var sum: u32 = 0;
    for (word) |c| {
        if (c == 10) {
            break;
        }
        sum += c;
        sum *= 17;
        sum = sum % 256;
    }

    return sum;
}

fn printBoxes(boxes: std.AutoHashMap(u32, std.ArrayList(Lens))) void {
    var it = boxes.iterator();
    while (it.next()) |list| {
        std.debug.print("Box {d}: ", .{list.key_ptr.*});

        for (list.value_ptr.items) |lens| {
            std.debug.print("[{s} {d}] ", .{ lens.label, lens.focalLength });
        }
        std.debug.print("\n", .{});
    }
}