#!/usr/bin/env python
#
from __future__ import annotations

from collections import deque
from dataclasses import dataclass
from typing import Callable, Iterable


def main():
    activity_tree = ActivityTree()
    activity_tree.activities[0] = Activity(id=0, parent=0, text="*")
    activity_tree.activities[1] = Activity(id=1, parent=0, text="A")
    activity_tree.activities[2] = Activity(id=2, parent=1, text="B")
    activity_tree.activities[3] = Activity(id=3, parent=2, text="C")
    activity_tree.activities[4] = Activity(id=4, parent=2, text="D")
    activity_tree.activities[5] = Activity(id=5, parent=4, text="E")
    activity_tree.activities[6] = Activity(id=6, parent=4, text="F")
    activity_tree.activities[7] = Activity(id=7, parent=1, text="G")
    activity_tree.activities[8] = Activity(id=8, parent=1, text="H")
    activity_tree.activities[9] = Activity(id=9, parent=8, text="I")
    activity_tree.activities[10] = Activity(id=10, parent=9, text="J")

    match_iter = MatchIter(activity_tree, lambda act: act.id in (2, 5, 10))

    while (depth_and_node := match_iter.get_next()) is not None:
        depth, matched, node = depth_and_node
        print("\t".join((str(depth), str(matched), node.text)))

    # while True:
    #     print("\n\n\nvvvvvvvvvv\n\n")
    #     node = match_iter.get_next()
    #     if node is None:
    #         break
    #     print(node.text)
    #     print("\n\n")
    #     match_iter.print_stack()
    #     print("\n\n\n^^^^^^^^^^\n\n")


class ActivityTree:
    activities: dict[int, Activity]

    def __init__(self) -> None:
        self.activities = {}


@dataclass
class Activity:
    id: int
    parent: int
    text: str


class MatchIter:
    activity_tree: ActivityTree
    stack: deque[StackEntry]
    is_match: Callable[[Activity], bool]

    def __init__(
        self, activity_tree: ActivityTree, is_match: Callable[[Activity], bool]
    ) -> None:
        self.activity_tree = activity_tree
        self.stack = deque([StackEntry(0)])
        self.is_match = is_match

    def get_next(self) -> tuple[int, bool, Activity] | None:
        if (next_announce := self.get_next_announce()) is not None:
            next_announce_depth, next_announce_entry = next_announce
            next_announce_entry.announced = True
            return (
                next_announce_depth,
                next_announce_entry.matched,
                self.activity_tree.activities[next_announce_entry.node_id],
            )

        while len(self.stack) > 0:
            stack_entry = self.stack[-1]
            node = self.activity_tree.activities[stack_entry.node_id]

            if stack_entry.visited:
                self.stack.pop()
                continue
            # print(f"Checking {stack_entry.text(self.activity_tree)}")

            if self.is_match(node):
                stack_entry.should_announce = True
                stack_entry.matched = True
                for parent in self.walk_up_stack(stack_entry):
                    if parent.node_id != 0 and not parent.should_announce:
                        parent.should_announce = True

            stack_entry.visited = True

            children = [
                c for c in self.activity_tree.activities.values() if c.parent == node.id
            ]
            # for c in children:
            for c in reversed(children):
                if c.id != 0:
                    self.stack.append(StackEntry(c.id))

            if (
                stack_entry.should_announce
                and (next_announce := self.get_next_announce()) is not None
            ):
                next_announce_depth, next_announce_entry = next_announce
                next_announce_entry.announced = True
                return (
                    next_announce_depth,
                    next_announce_entry.matched,
                    self.activity_tree.activities[next_announce_entry.node_id],
                )

        return None

    def print_stack(self) -> None:
        for stack_entry in self.stack:
            print(stack_entry.text(self.activity_tree))

    def get_next_announce(self) -> tuple[int, StackEntry] | None:
        for entry in self.stack:
            if entry.should_announce and not entry.announced:
                depth = len(list(self.walk_up_stack(entry)))

                return (depth, entry)
        return None

    def walk_up_stack(self, base: StackEntry) -> Iterable[StackEntry]:
        found_start = False
        for i in range(len(self.stack)):
            idx = -1 - i
            stack_entry = self.stack[idx]
            if not found_start:
                if stack_entry is base:
                    found_start = True
                continue
            if stack_entry.visited:
                yield stack_entry


@dataclass
class StackEntry:
    node_id: int
    announced: bool = False
    should_announce: bool = False
    visited: bool = False
    matched: bool = False

    def text(self, activity_tree: ActivityTree) -> str:
        node = activity_tree.activities[self.node_id]
        return f"{node.text}\t{"annced" if self.announced else "-"}\t{"should" if self.should_announce else "-"}\t{"visited" if self.visited else "-"}\t{"matched" if self.matched else "-"}"


if __name__ == "__main__":
    main()