from __future__ import annotations
from collections import defaultdict
from tools.aoc import AOCDay
from tools.math import mul
from typing import Any


class Part:
    def __init__(self, part_str: str):
        self.attr = defaultdict(int)

        for desc in part_str[1:-1].split(","):
            attr, value = desc.split("=")
            self.attr[attr] = int(value)

    def get_value(self) -> int:
        return sum(self.attr.values())


class WorkflowItem:
    def __init__(self, flow_str: str):
        if ":" in flow_str:
            check, self.check_target = flow_str.split(":")
            self.check_attr = check[0]
            self.check_comp = check[1]
            self.check_value = int(check[2:])
        else:
            self.check_target = flow_str
            self.check_comp = None
            self.check_attr = None
            self.check_value = None

    def get_target(self, part: Part) -> str | None:
        if self.check_attr is None:
            return self.check_target
        else:
            match self.check_comp:
                case "<":
                    result = part.attr[self.check_attr] < self.check_value
                case ">":
                    result = part.attr[self.check_attr] > self.check_value
                case _:
                    assert False, "unhandled comparator: %s" % self.check_comp

            if result:
                return self.check_target

    def reverse(self) -> WorkflowItem:
        if self.check_attr is None:
            return self
        else:
            if self.check_comp == ">":
                return WorkflowItem(f"{self.check_attr}<{self.check_value + 1}:{self.check_target}")
            else:
                return WorkflowItem(f"{self.check_attr}>{self.check_value - 1}:{self.check_target}")


class Workflow:
    def __init__(self, workflow_str: str):
        workflow_str = workflow_str[:-1]
        self.name, workflow = workflow_str.split("{")
        self.flow = [WorkflowItem(x) for x in workflow.split(",")]
        self.source = None

    def get_target(self, part: Part) -> str:
        for flow in self.flow:
            target = flow.get_target(part)
            if target is not None:
                return target

    def has_target(self, target: str) -> bool:
        return target in [x.check_target for x in self.flow]

    def reached_workflow(self, target: str) -> list[list[WorkflowItem]]:
        filter_lists = []
        my_list = []
        for item in self.flow:
            if item.check_target != target:
                my_list.append(item.reverse())
            else:
                if item.check_attr is None:
                    filter_lists.append(my_list)
                else:
                    filter_lists.append(my_list + [item])
                    my_list.append(item.reverse())

        if self.source is not None:
            new_filter_lists = []
            sub_lists = self.source.reached_workflow(self.name)
            for sub_list in sub_lists:
                for my_list in filter_lists:
                    new_filter_lists.append(sub_list + my_list)

            filter_lists = new_filter_lists

        return filter_lists


def get_unfiltered_parts(filters: list[WorkflowItem]) -> int:
    attr_dict = {
        "x": {">": 0, "<": 4001},
        "m": {">": 0, "<": 4001},
        "a": {">": 0, "<": 4001},
        "s": {">": 0, "<": 4001},
    }
    for item in filters:
        match item.check_comp:
            case ">":
                attr_dict[item.check_attr][">"] = max(attr_dict[item.check_attr][">"], item.check_value)
            case "<":
                attr_dict[item.check_attr]["<"] = min(attr_dict[item.check_attr]["<"], item.check_value)
            case _:
                assert False, f"unexpected check comp {item.check_comp}, {item.check_target}"

    return mul([v["<"] - v[">"] - 1 for v in attr_dict.values()])


class Day(AOCDay):
    inputs = [
        [
            (19114, "input19_test"),
            (397643, "input19"),
        ],
        [
            (167409079868000, "input19_test"),
            (132392981697081, "input19"),
        ],
    ]

    def parse_input(self) -> (dict[str, Workflow], list[Part]):
        workflow_list, parts = self.getMultiLineInputAsArray()
        workflow_list = [Workflow(x) for x in workflow_list]
        parts = [Part(x) for x in parts]

        workflows = {}
        for w in workflow_list:
            workflows[w.name] = w
            ps = [x for x in workflow_list if x.has_target(w.name)]
            if ps:
                assert len(ps) == 1
                w.source = ps[0]

        return workflows, parts

    def part1(self) -> Any:
        workflows, parts = self.parse_input()
        ans = 0
        for part in parts:
            target = workflows["in"].get_target(part)
            while target not in ["A", "R"]:
                target = workflows[target].get_target(part)

            if target == "A":
                ans += part.get_value()

        return ans

    def part2(self) -> Any:
        workflows, _ = self.parse_input()
        ans = 0
        for workflow in workflows.values():
            if workflow.has_target("A"):
                filters = workflow.reached_workflow("A")
                ans += sum(get_unfiltered_parts(x) for x in filters)
        return ans


if __name__ == "__main__":
    day = Day(2023, 19)
    day.run(verbose=True)