const std = @import("std");

const Rule = struct { before: u32, after: u32 };

const Input = struct {
    rules: []Rule,
    updates: [][]u32,
};

pub fn parse(allocator: std.mem.Allocator, data: []const u8) !Input {
    var lines = std.mem.split(u8, data, "\n");
    var rules = std.ArrayList(Rule).init(allocator);
    while (lines.next()) |line| {
        if (line.len == 0) break;
        var parts = std.mem.split(u8, line, "|");
        try rules.append(.{
            .before = try std.fmt.parseInt(u32, parts.next() orelse return error.InvalidInput, 10),
            .after = try std.fmt.parseInt(u32, parts.next() orelse return error.InvalidInput, 10),
        });
    }
    var updates = std.ArrayList([]u32).init(allocator);
    while (lines.next()) |line| {
        if (line.len == 0) break;
        var parts = std.mem.split(u8, line, ",");
        var update = std.ArrayList(u32).init(allocator);
        while (parts.next()) |part| {
            try update.append(try std.fmt.parseInt(u32, part, 10));
        }
        std.debug.assert(update.items.len % 2 == 1);
        try updates.append(update.items);
    }
    return .{ .rules = rules.items, .updates = updates.items };
}

const test_input =
    \\47|53
    \\97|13
    \\97|61
    \\97|47
    \\75|29
    \\61|13
    \\75|53
    \\29|13
    \\97|29
    \\53|29
    \\61|53
    \\97|53
    \\61|29
    \\47|13
    \\75|47
    \\97|75
    \\47|61
    \\75|61
    \\47|29
    \\75|13
    \\53|13
    \\
    \\75,47,61,53,29
    \\97,61,53,29,13
    \\75,29,13
    \\75,97,47,61,53
    \\61,13,29
    \\97,13,75,29,47
    \\
;

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

    const output = try part1(arena.allocator(), try parse(arena.allocator(), test_input));
    std.debug.print("got {}\n", .{output});
    std.debug.assert(output == 143);
}

pub fn part1(_: std.mem.Allocator, input: Input) !u32 {
    var sum: u32 = 0;
    for (input.updates) |update| {
        const ok = for (input.rules) |rule| {
            const first = std.mem.indexOfScalar(u32, update, rule.before) orelse continue;
            const second = std.mem.indexOfScalar(u32, update, rule.after) orelse continue;
            if (first > second) break false;
        } else true;
        if (ok) {
            sum += update[update.len / 2];
        }
    }
    return sum;
}

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

    const output = try part2(arena.allocator(), try parse(arena.allocator(), test_input));
    std.debug.print("got {}\n", .{output});
    std.debug.assert(output == 123);
}

pub fn part2(_: std.mem.Allocator, input: Input) !u32 {
    var sum: u32 = 0;
    for (input.updates) |update| {
        var reordered_any = false;
        var reordered = true;
        while (reordered) {
            reordered = false;
            for (input.rules) |rule| {
                const first = std.mem.indexOfScalar(u32, update, rule.before) orelse continue;
                const second = std.mem.indexOfScalar(u32, update, rule.after) orelse continue;
                if (first > second) {
                    reordered = true;
                    reordered_any = true;
                    std.mem.swap(u32, &update[first], &update[second]);
                }
            }
        }
        if (reordered_any) sum += update[update.len / 2];
    }
    return sum;
}