#!/usr/bin/env python3

from typing import List, Callable


class BorNode:
    """Узел префиксного дерева."""
    def __init__(self):
        # Узлы-потомки текущего узла (ключ - символ какой-либо строки-паттерна,
        # значение - объект BorNode).
        self.children = {}
        # Совпавшие паттерны.
        self.matched_patterns = []
        # Суффиксная ссылка (ссылка на BorNode, соответствующий самому длинному
        # суффиксу).
        self.suffix_link = None


class AhoCorasickSearcher:
    def __init__(self, patterns: List[str]):
        self._root_node = self._create_state_machine(patterns)

    @staticmethod
    def _build_prefix_tree(patterns: List[str]):
        """Создает бор - префиксное дерево из паттернов поиска."""
        root = BorNode()

        for path in patterns:
            node = root
            for symbol in path:
                if symbol not in node.children:
                    node.children[symbol] = BorNode()
                    node = node.children[symbol]
            # Для каждого листка дерева, в который прийдет автомат, он будет
            # в конечном состоянии, что будет означать, что данный паттерн path
            # присутствует в строке, по который происходил поиск.
            # Поэтому нужно добавить в этот список паттернов листка паттерн path.
            node.matched_patterns.append(path)
        return root

    def _create_state_machine(self, patterns: List[str]):
        """
        Создает конечный автомат: для всех узлов дерева находит и присваивает
        суффиксные ссылки.
        """
        # Создаем бор, инициализируем непосредственных потомков корневого узла.
        root = self._build_prefix_tree(patterns)
        queue = []
        for node in root.children.values():
            queue.append(node)
            node.fail = root

        # Инициализируем остальные узлы:
        # 1. Берем очередной узел (важно, что проход в ширину)
        # 2. Находим самую длинную суффиксную ссылку для этой вершины - это и будет fail-функция
        # 3. Если таковой не нашлось - устанавливаем fail-функцию в корневой узел

        # Идет поиск вширь по префиксному дереву.
        while len(queue) > 0:
            node = queue.pop(0)

            for symbol, child_node in node.children.items():
                # Добавляем очередной узел-потомок в очередь для корректного
                # дальнейшего поиска вширь по дереву.
                queue.append(child_node)
                # Сначала начинает с суфиксной ссылки текущего узла node.
                # Эта ссылка ведет на узел, описывающий состояние автомата, в
                # которое он перейдет, если прекратиться matching в строке.
                fail_node = node.suffix_link
                # Поиск самой длинной суффиксной ссылки для узла-потомка.
                while fail_node is not None and symbol not in fail_node.children:
                    fail_node = fail_node.suffix_link
                if fail_node is not None:
                    child_node.suffix_link = child_node.children[symbol]
                else:
                    child_node.suffix_link = root

                # Добавляем в текущий узел-потомок паттерны из узла по
                # суффиксной ссылке. Это нужно, чтобы не пропустить паттерны
                # которые бы мы получили если бы пошли по ветке дерева, которая
                # привела бы нас в child_node.suffix_link сразу.
                child_node.matched_patterns.extend(
                    child_node.suffix_link.matched_patterns
                )
        # Возвращаем ссылку на дерево-автомат.
        return root

    def find_all(self, target_string, matching_handler: Callable[[int, List[str]], None]):
        """
        Находит все возможные подстроки из набора паттернов в строке.

        Args:
            target_string: строка, по которой происходит поиск паттернов.
            matching_handler: обработчик совпадения какой-либо части строки
                target_string с некоторыми паттернами.
        """
        node = self._root_node

        for i, symbol in enumerate(target_string):
            # Если в какой-то момент при matching'е мы пришли в "неудачный" узел,
            # т.е. в узел, на котором matching прекратился, то нужно перевести
            # автомат в нужное состояние - перейти по суффиксной ссылке.
            while node is not None and symbol not in node.children:
                node = node.fail
            # Если очередная итерация matching'а закончилась в корневом узле,
            # то назначаем его текущему узлу для последующих итераций.
            if node is None:
                node = self._root_node
                continue
            # Если выполнение дошло до этого места, symbol ведет из текущего
            # узла node в другой, т.е. итерация матчинга завершилась "успешно".
            # Тогда запоминаем следующий узел, как текущий. И вызываем функцию-
            # обработчик со всеми паттернами, которые сохранены в нем (если,
            # конечно, этот узел терминальный и там паттерны вообще есть).
            node = node.children[symbol]
            for pattern in node.matched_patterns:
                matching_handler(i - len(pattern) + 1, pattern)


def _test():
    def on_occurence(pos, patterns):
        """
        Обработчик-заглушка, просто выводит позицию pos (индекс) на котором в
        очередной строки произошло совпадение по паттернам patterns.
        """
        print("At pos {} found pattern: {}".format(pos, patterns))

    patterns = ['a', 'ab', 'abc', 'bc', 'c', 'cba']
    target_string = "abcba"
    searcher = AhoCorasickSearcher(patterns)
    searcher.find_all(target_string, on_occurence)


_test()