const std = @import("std");
const mecha = @import("mecha");
const print = std.debug.print;

pub fn solve(part: []u8, buffer: []u8, allocator: std.mem.Allocator) !void {
    var lines = std.mem.splitScalar(u8, buffer, '\n');
    if (std.mem.eql(u8, part, "1")) {
        try part1(&lines, allocator);
    } else {
        try part2(&lines, allocator);
    }
}

fn get_ordering_rules(lines: *std.mem.SplitIterator(u8, .scalar), allocator: std.mem.Allocator) !std.AutoHashMap(usize, std.ArrayList(usize)) {
    var ordering_rules = std.AutoHashMap(usize, std.ArrayList(usize)).init(allocator);
    errdefer ordering_rules.deinit();

    const parser = mecha.combine(.{ mecha.int(usize, .{}), mecha.ascii.char('|').discard(), mecha.int(usize, .{}) });

    // PAGE ORDERING RULES
    while (lines.next()) |line| {
        if (line.len == 0) {
            break;
        }

        const parsed = try parser.parse(allocator, line);

        if (ordering_rules.getPtr(parsed.value[0])) |value| {
            try value.append(parsed.value[1]);
        } else {
            var list = std.ArrayList(usize).init(allocator);
            errdefer list.deinit();
            try list.append(parsed.value[1]);
            try ordering_rules.put(parsed.value[0], list);
        }
    }

    return ordering_rules;
}

fn part1(lines: *std.mem.SplitIterator(u8, .scalar), allocator: std.mem.Allocator) !void {
    var ordering_rules = try get_ordering_rules(lines, allocator);
    defer ordering_rules.deinit();

    var total: usize = 0;

    // UPDATES
    outer: while (lines.next()) |line| {
        if (line.len == 0) {
            break;
        }

        var list = std.ArrayList(usize).init(allocator);
        defer list.deinit();

        var covered_numbers = std.AutoHashMap(usize, void).init(allocator);
        defer covered_numbers.deinit();

        var page_numbers = std.mem.splitScalar(u8, line, ',');
        while (page_numbers.next()) |page_number| {
            const number = try std.fmt.parseInt(usize, page_number, 10);
            try covered_numbers.put(number, {});
            try list.append(number);

            const next = ordering_rules.get(number) orelse continue;
            for (next.items) |n| {
                if (covered_numbers.contains(n)) {
                    continue :outer;
                }
            }
        }

        total += list.items[list.items.len / 2];
    }

    print("{d}\n", .{total});
}

//-----------------------PART 2----------------------------

fn reorder(list: *std.ArrayList(usize), ordering_rules: *std.AutoHashMap(usize, std.ArrayList(usize)), allocator: std.mem.Allocator) !bool {
    var covered_numbers = std.AutoHashMap(usize, usize).init(allocator);
    defer covered_numbers.deinit();

    for (list.items, 0..) |number, i| {
        try covered_numbers.put(number, i);

        const next = ordering_rules.get(number) orelse continue;
        for (next.items) |n| {
            if (covered_numbers.get(n)) |index| {
                list.items[i] = n;
                list.items[index] = number;
                return false;
            }
        }
    }
    return true;
}

fn part2(lines: *std.mem.SplitIterator(u8, .scalar), allocator: std.mem.Allocator) !void {
    var ordering_rules = try get_ordering_rules(lines, allocator);
    defer ordering_rules.deinit();

    var total: usize = 0;

    // UPDATES
    while (lines.next()) |line| {
        if (line.len == 0) {
            break;
        }

        var list = std.ArrayList(usize).init(allocator);
        defer list.deinit();

        var page_numbers = std.mem.splitScalar(u8, line, ',');
        while (page_numbers.next()) |page_number| {
            const number = try std.fmt.parseInt(usize, page_number, 10);
            try list.append(number);
        }

        var incorrect = false;
        while (!try reorder(&list, &ordering_rules, allocator)) {
            incorrect = true;
        }

        if (incorrect) {
            total += list.items[list.items.len / 2];
        }
    }

    print("{d}\n", .{total});
}