from textwrap import dedent
from typing import Iterator, override
from unittest import TestCase

from puzzles._solver import Solver

Update = tuple[int, ...]


class DayFiveSolver(Solver):
    ordering_rules: dict[int, set[int]]
    updates: list[Update]

    @override
    def __init__(self, puzzle_input: str):
        ordering_rules_input, _, updates_input = puzzle_input.partition("\n\n")
        self.ordering_rules = {}

        for line in ordering_rules_input.strip().split("\n"):
            a, _, b = line.partition("|")
            a, b = int(a), int(b)
            if b not in self.ordering_rules:
                self.ordering_rules[b] = set()
            self.ordering_rules[b].add(a)

        self.updates = [
            tuple(map(int, line.split(",")))
            for line in updates_input.strip().split("\n")
        ]

    @override
    def solve_p1(self) -> int:
        return sum(
            update[len(update) // 2]
            for update in self.updates
            if self.is_correctly_ordered(update)
        )

    @override
    def solve_p2(self) -> int:
        return sum(
            self.fix_update(update)[len(update) // 2]
            for update in self.updates
            if not self.is_correctly_ordered(update)
        )

    def is_correctly_ordered(self, update: Update) -> bool:
        contains = set(update)
        seen = set()

        for i in update:
            for j in self.ordering_rules.get(i, ()):
                if j in contains and j not in seen:
                    return False
            seen.add(i)

        return True

    def fix_update(self, update: Update) -> Update:
        contains = set(update)
        seen = set()

        fixed_update = list()

        def fix_item(i):
            for j in self.ordering_rules.get(i, ()):
                if j in contains and j not in seen:
                    fix_item(j)
            fixed_update.append(i)
            seen.add(i)

        for i in update:
            if i not in seen:
                fix_item(i)

        return tuple(fixed_update)


class TestDayFiveSolver(TestCase):
    def test(self):
        solver = DayFiveSolver(
            dedent(
                """
                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
                """
            )
        )
        self.assertEqual(solver.solve_p1(), 143)
        self.assertEqual(solver.solve_p2(), 123)