#!/usr/bin/env python3

# https://adventofcode.com/2024/day/5

def get_lines(filename: str) -> list:
    with open(filename, "r") as file:
        return [line.strip() for line in file.readlines()]


def get_rules_and_updates(lines: list) -> tuple:
    rules = {}
    updates = []
    for line in lines:
        if "|" in line:
            a, b = line.split("|")
            if rules.get(b):
                rules[b].append(a)
            else:
                rules[b] = [a]
        elif "," in line:
            updates.append(line.split(","))
    return rules, updates


def check_validity(rules: dict, update: list) -> bool:
    for position, page_number in enumerate(update):
        if page_number not in rules:
            continue
        for earlier_number in rules[page_number]:
            if earlier_number in update[position:]:
                return False
    return True


def sort_update(rules: dict, update: list) -> None:
    len_update = len(update)
    for i in range(len_update - 1):
        for j in range(len_update - i - 1):
            if update[j] not in rules:
                continue
            if update[j + 1] in rules[update[j]]:
                update[j], update[j + 1] = update[j + 1], update[j]


def get_middle_page(update: list) -> int:
    middle_pos = len(update) // 2
    return int(update[middle_pos])


def get_middle_page_sums(rules: dict, updates: list) -> tuple:
    valid_middle_page_sum = 0
    corrected_middle_page_sum = 0
    for update in updates:
        if check_validity(rules, update):
            valid_middle_page_sum += get_middle_page(update)
        else:
            sort_update(rules, update)
            corrected_middle_page_sum += get_middle_page(update)
    return valid_middle_page_sum, corrected_middle_page_sum


def main():
    # lines = get_lines("sample-input.txt")
    lines = get_lines("input.txt")

    rules, updates = get_rules_and_updates(lines)
    valid_middle_page_sum, invalid_middle_page_sum = get_middle_page_sums(rules, updates)
    print("Part 1: The sum of middle pages in valid updates is:", valid_middle_page_sum)
    print("Part 2: The sum of middle pages in corrected updates is:", invalid_middle_page_sum)


if __name__ == '__main__':
    main()