const std = @import("std");

const PATH = "input/day07.txt";

const Str = []const u8;

/// Find every parent and parent of parents for `shiny gold`.
pub fn first(allocator: ?std.mem.Allocator) anyerror!usize {
    var arena = std.heap.ArenaAllocator.init(allocator.?);
    defer arena.deinit();

    var parents = try parseParents(arena.allocator());
    defer parents.deinit();

    var visited = std.StringHashMap(void).init(arena.allocator());
    defer visited.deinit();

    return try getParents(parents, &visited, "shiny gold");
}

fn getParents(parents: Parents, visited: *std.StringHashMap(void), bag: Str) anyerror!usize {
    var prts = parents.get(bag) orelse return 0;

    var sum: usize = 0;

    var it = prts.keyIterator();
    while (it.next()) |parent| {
        if (visited.contains(parent.*)) continue;
        try visited.put(parent.*, {});

        sum += 1;

        sum += try getParents(parents, visited, parent.*);
    }

    return sum;
}

const Parents = std.StringHashMap(std.StringHashMap(void));

/// parseParents collects the parents of each bag color. Each bag could have multiple
/// parents.
fn parseParents(allocator: std.mem.Allocator) !Parents {
    const file = @embedFile(PATH);
    var lines = std.mem.split(u8, file, "\n");

    var parents = Parents.init(allocator);

    while (lines.next()) |line| {
        if (line.len == 0) break;
        var par_chd = std.mem.split(u8, line, " bags contain ");
        const parent = par_chd.next().?;
        // std.debug.print("parent: {s}\n", .{parent});

        var children = std.mem.split(u8, par_chd.next().?, ", ");

        while (children.next()) |chd| {
            var chd_part = std.mem.tokenize(u8, chd, " ");

            const db = chd_part.next().?;
            if (std.mem.eql(u8, db, "no")) continue;

            const child_name = try std.mem.concat(allocator, u8, &[_]Str{ chd_part.next().?, " ", chd_part.next().? });
            // std.debug.print("child: {s}\n", .{child_name});

            var res = try parents.getOrPut(child_name);
            if (res.found_existing) try res.value_ptr.put(parent, {}) else {
                res.value_ptr.* = std.StringHashMap(void).init(allocator);
                try res.value_ptr.put(parent, {});
            }
        }
    }

    return parents;
}

/// Finds children and children of children of `shiny gold` bags.
pub fn second(allocator: ?std.mem.Allocator) anyerror!usize {
    var arena = std.heap.ArenaAllocator.init(allocator.?);
    defer arena.deinit();

    var children = try parseChildren(arena.allocator());
    defer children.deinit();

    return try getChildren(children, "shiny gold");
}

fn getChildren(children: Children, bag: Str) anyerror!usize {
    var kids = children.get(bag) orelse return 0;

    var sum: usize = 0;

    for (kids.items) |kid| {
        sum += kid.piece;
        sum += kid.piece * try getChildren(children, kid.kind);
    }

    return sum;
}

const ChildBag = struct {
    piece: usize,
    kind: Str,
};

const Children = std.StringHashMap(std.ArrayList(ChildBag));

/// parseChildren collects the child bag quantity and kind for every parent bag.
fn parseChildren(allocator: std.mem.Allocator) !Children {
    const file = @embedFile(PATH);
    var lines = std.mem.split(u8, file, "\n");

    var child_hash = Children.init(allocator);

    while (lines.next()) |line| {
        if (line.len == 0) break;
        var par_chd = std.mem.split(u8, line, " bags contain ");
        const parent = par_chd.next().?;

        var children = std.mem.split(u8, par_chd.next().?, ", ");

        while (children.next()) |chd| {
            var chd_part = std.mem.tokenize(u8, chd, " ");

            const db = chd_part.next().?;
            if (std.mem.eql(u8, db, "no")) continue;
            const db_num = try std.fmt.parseUnsigned(usize, db, 0);

            const child_name = try std.mem.concat(allocator, u8, &[_]Str{ chd_part.next().?, " ", chd_part.next().? });

            var res = try child_hash.getOrPut(parent);
            if (res.found_existing) try res.value_ptr.append(.{ .piece = db_num, .kind = child_name }) else {
                res.value_ptr.* = std.ArrayList(ChildBag).init(allocator);
                try res.value_ptr.append(.{ .piece = db_num, .kind = child_name });
            }
        }
    }

    return child_hash;
}

const test_input =
    \\light red bags contain 1 bright white bag, 2 muted yellow bags.
    \\dark orange bags contain 3 bright white bags, 4 muted yellow bags.
    \\bright white bags contain 1 shiny gold bag.
    \\muted yellow bags contain 2 shiny gold bags, 9 faded blue bags.
    \\shiny gold bags contain 1 dark olive bag, 2 vibrant plum bags.
    \\dark olive bags contain 3 faded blue bags, 4 dotted black bags.
    \\vibrant plum bags contain 5 faded blue bags, 6 dotted black bags.
    \\faded blue bags contain no other bags.
    \\dotted black bags contain no other bags.
;

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

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