import logging
import queue
from abc import ABC, abstractmethod
from collections import deque
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime
from dataclasses import dataclass
from threading import Thread, Event
from typing import TypeVar, Generic, Dict, Optional, List
import re
import time
from typing import Any

try:
    from aabd.base.patched_logging import get_logger
except:
    from logging import getLogger as get_logger

T = TypeVar('T')


@dataclass
class Message(Generic[T]):
    seq: int
    timestamp: datetime
    type: str
    header: dict
    payload: T


class MessageMaker:
    def __init__(self):
        self.seq = 0

    def data_message(self, header=None, payload=None):
        if header is None:
            header = {}
        if payload is None:
            payload = {}
        message = Message(seq=self.seq, timestamp=datetime.now(), header=header, type='data', payload=payload)
        self.seq += 1
        return message

    def stop_message(self):
        message = Message(seq=self.seq, timestamp=datetime.now(), type='stop', header={}, payload=None)
        self.seq += 1
        return message


from functools import wraps


def time_decorator(func):
    """装饰器：记录方法执行时间并打印信息"""

    @wraps(func)
    def wrapper(*args, **kwargs):
        start_time = time.time()
        try:
            result = func(*args, **kwargs)
        finally:
            elapsed = time.time() - start_time
            # args[0] 是 self（假设用于类方法）
            self_obj = args[0]
            self_obj.logger.debug(f"NODE_USE_TIME {self_obj.name} use time: {elapsed:.4f}s")
        return result

    return wrapper


class NodeMeta(type(ABC)):
    """元类：在子类定义时自动装饰 input 方法"""

    def __new__(cls, name, bases, attrs):
        new_class = super().__new__(cls, name, bases, attrs)
        # 检查子类是否实现了 input 方法
        if 'process' in attrs:
            original_method = attrs['process']
            # 应用装饰器
            decorated_method = time_decorator(original_method)
            setattr(new_class, 'process', decorated_method)
        return new_class


class Node(ABC, metaclass=NodeMeta):

    def __init__(self, context=None, name=None, logger=None, **kwargs):
        self.logger_root_name = 'node'
        self.name = name or self.__class__.__name__
        self.description = self.name
        self.output_nodes = set()
        self.input_nodes = set()
        self.context = context
        self.set_logger(logger)
        self.exception_callback = lambda key, node, message, e: None
        self.event_bus = None
        super().__init__(**kwargs)

    @abstractmethod
    def input(self, message, from_node: str):
        pass

    def set_logger(self, logger):
        if logger is None:
            logger = get_logger(__name__)
            self.logger_root_name = logger.name
        elif isinstance(logger, str):
            self.logger_root_name = logger
            logger = get_logger(f"{logger}.nodes")
        elif isinstance(logger, logging.Logger):
            self.logger_root_name = logger.name
            logger = get_logger(f"{logger.name}.nodes")
        else:
            self.logger_root_name = logger.name
            logger = logger
        self.logger = logger

    def get_logger_by_name(self, name):
        return get_logger(f"{self.logger_root_name}.{name}")

    def registry_output(self, node):
        node.registry_input(self)
        self.output_nodes.add(node)
        return self

    def registry_input(self, node):
        self.input_nodes.add(node)

    def set_context(self, context):
        self.context = context


class SingleInputNode(Node):
    def __init__(self, context=None, name=None, enable=True, logger=None, **kwargs):
        name = name or self.__class__.__name__
        super().__init__(context, name, logger=logger)
        self.enable = enable

    def input(self, message: Message, from_node: str):
        if self.enable is False:
            for output_node in self.output_nodes:
                output_node.input(message, self.name)
            return
        if message.type == 'stop':
            self.stop()
            for out_node in self.output_nodes:
                self.logger.debug(
                    f'NODE_DATA_STREAM [{self.name}-->{out_node.name}] {message.seq:06d} {message.timestamp}')
                try:
                    out_node.input(message, self.name)
                except Exception as e:
                    self.exception_callback(f'NODE_{self.name}', self, message, e)
                    self.logger.exception(f'{self.name} process error')
        else:
            try:
                self.process([message])
            except Exception as e:
                self.exception_callback(f'NODE_{self.name}', self, message, e)
                self.logger.exception(f'{self.name} process error')
            for out_node in self.output_nodes:
                self.logger.debug(
                    f'NODE_DATA_STREAM [{self.name}-->{out_node.name}] {message.seq:06d} {message.timestamp}')
                try:
                    out_node.input(message, self.name)
                except Exception as e:
                    self.exception_callback(f'NODE_{self.name}', self, message, e)
                    self.logger.exception(f'{self.name} process error')
                self.logger.debug(
                    f'NODE_DATA_STREAM END [{self.name}-->{out_node.name}] {message.seq:06d} {message.timestamp}')

    @abstractmethod
    def process(self, messages: list[Message]):
        pass

    def stop(self):
        pass


class WithHistoryMessageNode(Node):
    def __init__(self, context=None, name=None, history_len=1, enable=True, logger=None, **kwargs):
        name = name or self.__class__.__name__
        super().__init__(context, name, logger=logger)
        self.history_len = history_len
        self.history_queue = deque(maxlen=history_len)
        self.enable = enable

    def input(self, message: Message, from_node: str):
        if self.enable is False:
            for output_node in self.output_nodes:
                output_node.input(message, self.name)
            return
        if message.type == 'stop':
            self.stop()
            self.history_queue.clear()
            for out_node in self.output_nodes:
                self.logger.debug(
                    f'NODE_DATA_STREAM [{self.name}-->{out_node.name}] {message.seq:06d} {message.timestamp}')
                try:
                    out_node.input(message, self.name)
                except Exception as e:
                    self.exception_callback(f'NODE_{self.name}', self, message, e)
                    self.logger.exception(f'{self.name} process error')
        else:
            self.process([message])
            self.history_queue.append(message)
            for out_node in self.output_nodes:
                self.logger.debug(
                    f'NODE_DATA_STREAM [{self.name}-->{out_node.name}] {message.seq:06d} {message.timestamp}')
                try:
                    out_node.input(message, self.name)
                except Exception as e:
                    self.exception_callback(f'NODE_{self.name}', self, message, e)
                    self.logger.exception(f'{self.name} process error')

    @abstractmethod
    def process(self, messages: list[Message]):
        pass

    @property
    def history_messages(self):
        return list(self.history_queue)

    def stop(self):
        pass


class BoundedThreadPool:
    def __init__(self, max_workers: int, max_queue_size: int):
        import threading
        self.task_queue = queue.Queue(maxsize=max_queue_size)
        self.threads = []
        self._shutdown = False

        for _ in range(max_workers):
            t = threading.Thread(target=self._worker, daemon=True)
            t.start()
            self.threads.append(t)

    def _worker(self):
        while True:
            try:
                task = self.task_queue.get(timeout=0.5)  # 避免永久阻塞，便于 shutdown
                if task is None:
                    break
                func, args, kwargs, future = task
                try:
                    result = func(*args, **kwargs)
                    future.set_result(result)
                except Exception as e:
                    future.set_exception(e)
                finally:
                    self.task_queue.task_done()
            except queue.Empty:
                continue

    def submit(self, fn, *args, **kwargs):
        from concurrent.futures import Future
        if self._shutdown:
            raise RuntimeError("Cannot submit after shutdown")
        future = Future()
        # 这里会阻塞，直到队列有空位！
        self.task_queue.put((fn, args, kwargs, future))
        return future

    def shutdown(self, wait=True):
        self._shutdown = True
        if wait:
            for _ in self.threads:
                self.task_queue.put(None)  # poison pill
            for t in self.threads:
                t.join()


class ThreadSingleInputNode(Node):
    def __init__(self, context=None, name=None, pool_size=5, max_workers=5, logger=None, **kwargs):
        name = name or self.__class__.__name__
        super().__init__(context, name, logger=logger)
        self.thread_pool = BoundedThreadPool(max_workers=max_workers, max_queue_size=pool_size)
        # self.thread_pool = ThreadPoolExecutor(max_workers=max_workers)
        # self._work_queue = queue.Queue(maxsize=pool_size)
        # self.thread_pool._work_queue = self._work_queue
        self.last_future = None

    def input(self, message: Message, from_node: str):
        if message.type == 'stop':
            self.thread_pool.shutdown(wait=True)
            for out_node in self.output_nodes:
                self.logger.debug(
                    f'NODE_DATA_STREAM [{self.name}-->{out_node.name}] {message.seq:06d} {message.timestamp}')
                try:
                    out_node.input(message, self.name)
                except Exception as e:
                    self.exception_callback(f'NODE_{self.name}', self, message, e)
                    self.logger.exception(f'{self.name} process error')
                self.logger.debug(
                    f'NODE_DATA_STREAM END [{self.name}-->{out_node.name}] {message.seq:06d} {message.timestamp}')
        else:
            def post_process():
                for out_node in self.output_nodes:
                    self.logger.debug(
                        f'NODE_DATA_STREAM [{self.name}-->{out_node.name}] {message.seq:06d} {message.timestamp}')
                    try:
                        out_node.input(message, self.name)
                    except Exception as e:
                        self.exception_callback(f'NODE_{self.name}', self, message, e)
                        self.logger.exception(f'{self.name} process error')
                    self.logger.debug(
                        f'NODE_DATA_STREAM END [{self.name}-->{out_node.name}] {message.seq:06d} {message.timestamp}')

            self.last_future = self.thread_pool.submit(self.thread_process, message, self.last_future, post_process)

    def thread_process(self, message: Message, future, func):
        try:
            self.process([message])
        except Exception as e:
            self.exception_callback(f'NODE_{self.name}', self, message, e)
            self.logger.exception('')
        if future is not None:
            future.result()
        func()

    @abstractmethod
    def process(self, messages: list[Message]):
        pass


class QueueNode(Node):
    def __init__(self, queue_size=1, context=None, name=None, logger=None, **kwargs):
        super().__init__(context, name, logger=logger)
        self.input_queue = queue.Queue(maxsize=queue_size)
        self.output_queue = queue.Queue(maxsize=queue_size)
        Thread(target=self._run_input_thread, daemon=True).start()
        Thread(target=self._run_output_thread, daemon=True).start()

    def process(self, messages: list[Message]):
        pass

    def worker_loop_init(self):
        pass

    def _run_input_thread(self):
        # self.worker_loop_init()
        while True:
            message = None
            try:
                message = self.input_queue.get(timeout=0.5)
                if message.type == 'stop':
                    self.output_queue.put(message)
                    break
                else:
                    self.process([message])
                    self.output_queue.put(message)
            except queue.Empty:
                continue
            except Exception as e:
                self.exception_callback(f'NODE_{self.name}', self, message, e)
                self.logger.exception(f'{self.name} process error')

    def _run_output_thread(self):
        while True:
            message = None
            try:
                message = self.output_queue.get(timeout=0.5)
                if message.type == 'stop':
                    pass
                else:
                    self.process([message])
            except queue.Empty:
                continue
            except Exception as e:
                self.exception_callback(f'NODE_{self.name}', self, message, e)
                self.logger.exception(f'{self.name} process error')
            if message is not None:
                for out_node in self.output_nodes:
                    self.logger.debug(
                        f'NODE_DATA_STREAM [{self.name}-->{out_node.name}] {message.seq:06d} {message.timestamp}')
                    try:
                        out_node.input(message, self.name)
                    except Exception as e:
                        self.exception_callback(f'NODE_{self.name}', self, message, e)
                        self.logger.exception(f'{self.name} process error')

    def input(self, message, from_node: str):
        self.input_queue.put(message)


class MultipleInputNode(Node):

    def __init__(self, context=None, name=None, logger=None, **kwargs):
        name = name or self.__class__.__name__
        super().__init__(context, name, logger=logger)
        self.inputs_seqs_map = {}
        self.current_seq = 0
        self.messages_map = {}

    def input(self, message: Message, from_node: str):
        self.inputs_seqs_map[from_node].add(message.seq)
        self.messages_map[message.seq] = message
        if all(self.current_seq in input_seqs for input_seqs in self.inputs_seqs_map.values()):
            message_out = self.messages_map.pop(self.current_seq)
            try:
                self.process([message_out])
            except Exception as e:
                self.exception_callback(f'NODE_{self.name}', self, message, e)
                self.logger.exception(f'{self.name} process error')
            for out_node in self.output_nodes:
                self.logger.debug(
                    f'NODE_DATA_STREAM [{self.name}-->{out_node.name}] {message_out.seq:06d} {message_out.timestamp}')
                try:
                    out_node.input(message_out, self.name)
                except Exception as e:
                    self.exception_callback(f'NODE_{self.name}', self, message, e)
                    self.logger.exception(f'{self.name} process error')
            for input_seqs in self.inputs_seqs_map.values():
                input_seqs.remove(self.current_seq)
            self.current_seq += 1

    def registry_input(self, node):
        self.inputs_seqs_map[node.name] = set()
        super().registry_input(node)

    def process(self, messages: list[Message]):
        pass


class DelayedNode(Node):
    def __init__(self, context=None, name=None, delay_len=10, only_full_process=False, enable=True, logger=None,
                 **kwargs):
        name = name or self.__class__.__name__
        super().__init__(context, name, logger=logger)
        self.delay_len = delay_len
        self.delay_queue = deque(maxlen=delay_len)
        #  首次处理延迟
        self.only_full_process = only_full_process
        self.enable = enable

    def input(self, message: Message, from_node: str):
        if self.enable is False:
            for output_node in self.output_nodes:
                output_node.input(message, self.name)
            return
        if message.type == 'stop':
            try:
                self.stop(list(self.delay_queue))
            except Exception as e:
                self.exception_callback(f'NODE_{self.name}', self, message, e)
                self.logger.exception(f'{self.name} process error')
            while len(self.delay_queue) > 0:
                message_out = self.delay_queue.popleft()
                for output_node in self.output_nodes:
                    self.logger.debug(
                        f'NODE_DATA_STREAM [{self.name}-->{output_node.name}] {message_out.seq:06d} {message_out.timestamp}')
                    try:
                        output_node.input(message_out, self.name)
                    except Exception as e:
                        self.exception_callback(f'NODE_{self.name}', self, message, e)
                        self.logger.exception(f'{self.name} process error')
            # for output_node in self.output_nodes:
            #     try:
            #         output_node.input(message, self.name)
            #     except:
            #         self.logger.exception(f'{self.name} process error')
        else:
            self.delay_queue.append(message)
            try:
                if not self.only_full_process or len(self.delay_queue) == self.delay_len:
                    self.process(list(self.delay_queue))
            except Exception as e:
                self.exception_callback(f'NODE_{self.name}', self, message, e)
                self.logger.exception(f'{self.name} process error')
            if len(self.delay_queue) == self.delay_len:
                message_out = self.delay_queue.popleft()
                for output_node in self.output_nodes:
                    self.logger.debug(
                        f'NODE_DATA_STREAM [{self.name}-->{output_node.name}] {message_out.seq:06d} {message_out.timestamp}')
                    try:
                        output_node.input(message_out, self.name)
                    except Exception as e:
                        self.exception_callback(f'NODE_{self.name}', self, message, e)
                        self.logger.exception(f'{self.name} process error')

    def process(self, messages: list[Message]):
        pass

    def stop(self, messages: list[Message]):
        pass


class BatchNode(Node):

    def __init__(self, context=None, name=None, batch_size=10, logger=None, **kwargs):
        name = name or self.__class__.__name__
        super().__init__(context, name, logger=logger)
        self.batch_size = batch_size
        self.data = []

    def input(self, message, from_node: str):
        if message.type == 'stop':
            if len(self.data) > 0:
                try:
                    self.process(self.data)
                except Exception as e:
                    self.exception_callback(f'NODE_{self.name}', self, message, e)
                    self.logger.exception(f'{self.name} process error')
        else:
            self.data.append(message)
            if len(self.data) == self.batch_size:
                try:
                    self.process(self.data)
                except Exception as e:
                    self.exception_callback(f'NODE_{self.name}', self, message, e)
                    self.logger.exception(f'{self.name} process error')
            else:
                return

        for message_out in self.data:
            for output_node in self.output_nodes:
                self.logger.debug(
                    f'NODE_DATA_STREAM [{self.name}-->{output_node.name}] {message_out.seq:06d} {message_out.timestamp}')
                try:
                    output_node.input(message_out, self.name)
                except Exception as e:
                    self.exception_callback(f'NODE_{self.name}', self, message, e)
                    self.logger.exception(f'{self.name} process error')
        self.data = []

        if message.type == 'stop':
            for output_node in self.output_nodes:
                output_node.input(message, self.name)

    def process(self, messages: list[Message]):
        pass

    def stop(self):
        pass


class MultiFunctionNode(Node):

    def __init__(self, context=None, name=None, delay_len=1, history_len=1, only_full_process=False,
                 one_by_one_stop=False, enable=True, logger=None, **kwargs):
        name = name or self.__class__.__name__
        super().__init__(context, name, logger=logger)
        self.delay_len = delay_len
        self.delay_queue = deque(maxlen=delay_len)
        self.only_full_process = only_full_process
        self.history_len = history_len
        self.history_queue = deque(maxlen=history_len)
        self.one_by_one_stop = one_by_one_stop
        self.segment_flag = '-1'  # yes no
        self.enable = enable

    def input(self, message, from_node: str):
        if self.enable is False:
            for output_node in self.output_nodes:
                output_node.input(message, self.name)
            return
        if message.type == 'stop':
            try:
                if len(self.delay_queue) > 0:
                    if self.one_by_one_stop:
                        while len(self.delay_queue) > 0:
                            message_out = self.delay_queue.popleft()
                            for output_node in self.output_nodes:
                                self.logger.debug(
                                    f'NODE_DATA_STREAM [{self.name}-->{output_node.name}] {message_out.seq:06d} {message_out.timestamp}')
                                try:
                                    output_node.input(message_out, self.name)
                                    self.history_queue.append(message_out)
                                except Exception as e:
                                    self.exception_callback(f'NODE_{self.name}', self, message, e)
                                    self.logger.exception(f'{self.name} process error')
                    try:
                        self.segment_stop(list(self.delay_queue))
                    except Exception as e:
                        self.exception_callback(f'NODE_{self.name}', self, message, e)
                        self.logger.exception(f'{self.name} process error')
                self.stop(list(self.delay_queue))
                while len(self.delay_queue) > 0:
                    message_out = self.delay_queue.popleft()
                    for output_node in self.output_nodes:
                        self.logger.debug(
                            f'NODE_DATA_STREAM [{self.name}-->{output_node.name}] {message_out.seq:06d} {message_out.timestamp}')
                        try:
                            output_node.input(message_out, self.name)
                        except Exception as e:
                            self.exception_callback(f'NODE_{self.name}', self, message, e)
                            self.logger.exception(f'{self.name} process error')
            except Exception as e:
                self.exception_callback(f'NODE_{self.name}', self, message, e)
                self.logger.exception(f'{self.name} process error')
            while len(self.delay_queue) > 0:
                message_out = self.delay_queue.popleft()
                for output_node in self.output_nodes:
                    self.logger.debug(
                        f'NODE_DATA_STREAM [{self.name}-->{output_node.name}] {message_out.seq:06d} {message_out.timestamp}')
                    try:
                        output_node.input(message_out, self.name)
                        self.history_queue.append(message_out)
                    except Exception as e:
                        self.exception_callback(f'NODE_{self.name}', self, message, e)
                        self.logger.exception(f'{self.name} process error')
            self.history_queue.clear()
            for output_node in self.output_nodes:
                self.logger.debug(
                    f'NODE_DATA_STREAM [{self.name}-->{output_node.name}] {message.seq:06d} {message.timestamp}')
                try:
                    output_node.input(message, self.name)
                    self.history_queue.append(message)
                except Exception as e:
                    self.exception_callback(f'NODE_{self.name}', self, message, e)
                    self.logger.exception(f'{self.name} process error')
        else:
            if self.is_segment_message(message, None if len(self.history_messages) == 0 else self.history_messages[-1]):

                self.delay_queue.append(message)
                try:
                    if not self.only_full_process or len(self.delay_queue) == self.delay_len:
                        if self.segment_flag != 'yes':
                            self.segment_start()
                            self.segment_flag = 'yes'
                        out_number = self.process(list(self.delay_queue))
                        # 如果返回一个数字说明这些需要输出
                        if out_number is not None and out_number > 0:
                            for _ in range(out_number):
                                message_out = self.delay_queue.popleft()
                                for output_node in self.output_nodes:
                                    self.logger.debug(
                                        f'NODE_DATA_STREAM [{self.name}-->{output_node.name}] {message_out.seq:06d} {message_out.timestamp}')
                                    try:
                                        output_node.input(message_out, self.name)
                                        self.history_queue.append(message_out)
                                    except Exception as e:
                                        self.exception_callback(f'NODE_{self.name}', self, message, e)
                                        self.logger.exception(f'{self.name} process error')
                except Exception as e:
                    self.exception_callback(f'NODE_{self.name}', self, message, e)
                    self.logger.exception(f'{self.name} process error after Seq {message.seq} incoming')
                if len(self.delay_queue) == self.delay_len:
                    message_out = self.delay_queue.popleft()
                    for output_node in self.output_nodes:
                        self.logger.debug(
                            f'NODE_DATA_STREAM [{self.name}-->{output_node.name}] {message_out.seq:06d} {message_out.timestamp}')
                        try:
                            output_node.input(message_out, self.name)
                            self.history_queue.append(message_out)
                        except Exception as e:
                            self.exception_callback(f'NODE_{self.name}', self, message, e)
                            self.logger.exception(f'{self.name} process error')
            else:
                # 不符合条件的消息，直接输出
                # 如果有历史数据进行历史数据的处理
                if len(self.delay_queue) > 0:
                    if self.one_by_one_stop:
                        while len(self.delay_queue) > 0:
                            message_out = self.delay_queue.popleft()
                            for output_node in self.output_nodes:
                                self.logger.debug(
                                    f'NODE_DATA_STREAM [{self.name}-->{output_node.name}] {message_out.seq:06d} {message_out.timestamp}')
                                try:
                                    output_node.input(message_out, self.name)
                                    self.history_queue.append(message_out)
                                except Exception as e:
                                    self.exception_callback(f'NODE_{self.name}', self, message, e)
                                    self.logger.exception(f'{self.name} process error')
                    try:
                        self.segment_stop(list(self.delay_queue))
                    except Exception as e:
                        self.exception_callback(f'NODE_{self.name}', self, message, e)
                        self.logger.exception(f'{self.name} process error')
                # 不多余
                while len(self.delay_queue) > 0:
                    message_out = self.delay_queue.popleft()
                    for output_node in self.output_nodes:
                        self.logger.debug(
                            f'NODE_DATA_STREAM [{self.name}-->{output_node.name}] {message_out.seq:06d} {message_out.timestamp}')
                        try:
                            output_node.input(message_out, self.name)
                            self.history_queue.append(message_out)
                        except Exception as e:
                            self.exception_callback(f'NODE_{self.name}', self, message, e)
                            self.logger.exception(f'{self.name} process error')
                self.history_queue.clear()
                self.segment_flag = 'no'
                # 当前的消息直接送出
                for output_node in self.output_nodes:
                    self.logger.debug(
                        f'NODE_DATA_STREAM [{self.name}-->{output_node.name}] {message.seq:06d} {message.timestamp}')
                    try:
                        output_node.input(message, self.name)
                    except Exception as e:
                        self.exception_callback(f'NODE_{self.name}', self, message, e)
                        self.logger.exception(f'{self.name} process error')

    @abstractmethod
    def process(self, messages: list[Message]):
        pass

    def is_segment_message(self, message_new: Message, message_old: Message):
        return True

    def segment_start(self):
        pass

    def segment_stop(self, messages: list[Message]):
        pass

    def stop(self, messages: list[Message]):
        pass

    @property
    def history_messages(self):
        return list(self.history_queue)


import torch


class AsyncCudaStreamNode(Node):
    def __init__(self, context=None, name=None, device="cuda:0", logger=None, **kwargs):
        super().__init__(context, name, logger=logger)
        self.device = torch.device(device)
        self.stream = torch.cuda.Stream(device=self.device)
        self.completion_event = torch.cuda.Event()

        # 新增：记录每个上游节点对应的 event（用于多输入）
        self.upstream_events: Dict[str, torch.cuda.Event] = {}

        self.inputs_seqs_map = {}
        self.current_seq = 0
        self.messages_map = {}

    def registry_output(self, node: "Node"):
        """注册下游，并将本节点的 completion_event 传递给它（如果是 CUDA 节点）"""
        node.registry_input(self)

        # 👇 关键：如果下游也是 AsyncCudaStreamNode，告诉它“等我完成”
        if isinstance(node, AsyncCudaStreamNode):
            # 下游节点记录：来自 self.name 的输入需要等待 self.completion_event
            node._register_upstream_event(self.name, self.completion_event)

        self.output_nodes.add(node)
        return self

    def _register_upstream_event(self, from_node_name: str, event: torch.cuda.Event):
        """被下游调用：注册一个上游 event"""
        self.upstream_events[from_node_name] = event

    def input(self, message: Message, from_node: str):
        if message.type == 'stop':
            for out_node in self.output_nodes:
                self.logger.debug(
                    f'NODE_DATA_STREAM [{self.name}-->{out_node.name}] {message.seq:06d} {message.timestamp}')
                try:
                    out_node.input(message, self.name)
                except Exception as e:
                    self.exception_callback(f'NODE_{self.name}', self, message, e)
                    self.logger.exception(f'{self.name} process error')
        else:
            """
            统一入口：无论是 CPU 还是 GPU 上游，都走这里
            """
            # 如果有来自该上游的 event，就等待
            if from_node in self.upstream_events:
                event = self.upstream_events[from_node]
                self.stream.wait_event(event)  # 异步等待，不阻塞 CPU

            try:
                # 启动处理
                with torch.cuda.stream(self.stream):
                    self.process([message], stream=self.stream)

            except Exception as e:
                self.exception_callback(f'NODE_{self.name}', self, message, e)
                self.logger.exception(f'{self.name} process error')

            # 记录本节点完成事件
            self.completion_event.record(self.stream)

            for out_node in self.output_nodes:
                self.logger.debug(
                    f'NODE_DATA_STREAM [{self.name}-->{out_node.name}] {message.seq:06d} {message.timestamp}')
                try:
                    if isinstance(out_node, AsyncCudaStreamNode):
                        # GPU → GPU：直接 input，下游会查 upstream_events 并 wait
                        out_node.input(message, from_node=self.name)
                    else:
                        # GPU → CPU：必须同步
                        self.stream.synchronize()
                        out_node.input(message, from_node=self.name)
                except Exception as e:
                    self.exception_callback(f'NODE_{self.name}', self, message, e)
                    self.logger.exception(f'{self.name} process error')

    @abstractmethod
    def process(self, messages: list[Message], stream: torch.cuda.Stream):
        pass


class QueueCudaStreamNode(Node):
    def __init__(self, context=None, name=None, logger=None, device='cuda:0', maxsize=5, **kwargs):
        super().__init__(context, name, logger=logger)
        self.device = torch.device(device) if isinstance(device, str) else device
        # 工作队列：每个元素是 (message, from_node, upstream_event)
        self._work_queue = queue.Queue(maxsize=maxsize)
        # 启动工作线程
        self._worker_thread_init_event = Event()
        self._worker_thread = Thread(target=self._worker_loop, daemon=True)
        self._worker_thread.start()
        self._worker_thread_init_event.wait()

    def worker_loop_init(self):
        pass

    def message_pre_process(self, message: Message):
        return message

    def message_post_process(self, message: Message):
        return message

    def _worker_loop(self):
        try:
            torch.cuda.set_device(self.device)  # 线程局部设置设备
            self.worker_loop_init()
            stream = torch.cuda.Stream(device=self.device)
        finally:
            self._worker_thread_init_event.set()
        while True:
            # 从队列获取：消息 + 来源节点名 + 上游事件（可能为 None）
            message, from_node, upstream_event = self._work_queue.get()

            if message.type == 'stop':
                # 广播 stop 消息（无 GPU 依赖）
                for out_node in self.output_nodes:
                    out_node.input(message, from_node=self.name)
                self.stop()
                self._work_queue.task_done()
                break
            try:
                # ✅ 在工作线程中等待上游完成（如果存在）
                if upstream_event is not None:
                    stream.wait_event(upstream_event)

                # 执行计算
                with torch.cuda.stream(stream):
                    message = self.message_pre_process(message)
                    self.process([message], stream=stream)
                    message = self.message_post_process(message)

                # 记录本节点完成事件
                local_event = torch.cuda.Event()
                local_event.record(stream)

                # 通知下游
                self._notify_downstream(message, local_event)

            except Exception as e:
                self.exception_callback(f'NODE_{self.name}', self, message, e)
                if self.logger:
                    self.logger.exception(f'{self.name} process error')
            finally:
                self._work_queue.task_done()

        torch.cuda.empty_cache()

    def _notify_downstream(self, message: 'Message', local_event: torch.cuda.Event):
        """将消息和完成事件传递给所有下游节点"""
        for out_node in self.output_nodes:
            if self.logger:
                self.logger.debug(
                    f'NODE_DATA_STREAM [{self.name}-->{out_node.name}] '
                    f'{message.seq:06d} {message.timestamp}'
                )
            try:
                if isinstance(out_node, QueueCudaStreamNode):
                    # 下游是 CUDA 节点：传递事件，实现异步依赖
                    out_node.input(message, from_node=self.name, upstream_event=local_event)
                else:
                    # 下游是 CPU 节点：同步后传递
                    local_event.synchronize()
                    out_node.input(message, from_node=self.name)
            except Exception as e:
                self.exception_callback(f'NODE_{self.name}', self, message, e)
                if self.logger:
                    self.logger.exception(f'{self.name} notify downstream error')

    def stop(self):
        """停止所有工作线程"""
        pass

    def input(self, message: 'Message', from_node: str, upstream_event: Optional[torch.cuda.Event] = None):
        """
        主线程入口：快速入队，不执行任何 CUDA 操作。

        参数:
            message: 输入消息
            from_node: 消息来源节点名
            upstream_event: （可选）上游完成事件，用于建立 GPU 依赖
        """
        self._work_queue.put((message, from_node, upstream_event))

    def process(self, messages: List['Message'], stream: torch.cuda.Stream):
        """
        子类必须实现此方法。
        在指定 CUDA 流中处理消息。
        """
        raise NotImplementedError("Subclasses must implement the 'process' method.")


class QueueSingleNode(Node):
    def __init__(self, context=None, name=None, logger=None, maxsize=5, **kwargs):
        super().__init__(context, name, logger=logger)
        # 工作队列：每个元素是 (message, from_node, upstream_event)
        self._work_queue = queue.Queue(maxsize=maxsize)

        self._worker_thread_init_event = Event()
        self._worker_thread = Thread(target=self._worker_loop, daemon=True)
        self._worker_thread.start()
        self._worker_thread_init_event.wait()

    def worker_loop_init(self):
        pass

    def message_pre_process(self, message: 'Message'):
        return message

    def message_post_process(self, message: 'Message'):
        return message

    def _worker_loop(self):
        try:
            self.worker_loop_init()
        finally:
            self._worker_thread_init_event.set()

        while True:
            message, from_node = self._work_queue.get()

            if message.type == 'stop':
                for out_node in self.output_nodes:
                    out_node.input(message, from_node=self.name)
                self.stop()
                self._work_queue.task_done()
                break

            try:
                # 忽略 upstream_event（CPU 节点无需 GPU 同步）
                message = self.message_pre_process(message)
                self.process([message])  # 注意：无 stream 参数
                message = self.message_post_process(message)

                # 通知下游（不传递事件）
                self._notify_downstream(message)

            except Exception as e:
                self.exception_callback(f'NODE_{self.name}', self, message, e)
                if self.logger:
                    self.logger.exception(f'{self.name} process error')
            finally:
                self._work_queue.task_done()

        # 可选：清理（CPU 场景通常不需要）
        pass

    def _notify_downstream(self, message: 'Message'):
        """纯 CPU 通知，不传递事件"""
        for out_node in self.output_nodes:
            if self.logger:
                self.logger.debug(
                    f'NODE_DATA [{self.name}-->{out_node.name}] '
                    f'{message.seq:06d} {message.timestamp}'
                )
            try:
                # 兼容 CUDA 节点：但 CPU 节点不传 upstream_event
                # 因此调用时只传两个参数（符合 input 签名的默认行为）
                out_node.input(message, from_node=self.name)
            except Exception as e:
                self.exception_callback(f'NODE_{self.name}', self, message, e)
                if self.logger:
                    self.logger.exception(f'{self.name} notify downstream error')

    def stop(self):
        pass

    def input(self, message: 'Message', from_node: str):
        """
        与 QueueCudaStreamNode 接口完全一致。
        upstream_event 在 CPU 节点中被忽略。
        """
        self._work_queue.put((message, from_node))

    def process(self, messages: List['Message']):
        """
        子类必须实现。
        注意：无 stream 参数（区别于 CUDA 版本）。
        """
        raise NotImplementedError("Subclasses must implement the 'process' method.")


class CompositeQueueNode(Node):
    def __init__(self, node_type='single', context=None, name=None, logger=None, maxsize=5, device='cuda:0', **kwargs):
        super().__init__(context, name, logger=logger)
        if node_type == 'cuda':
            self.device = torch.device(device) if isinstance(device, str) else device
            self._worker_thread = Thread(target=self._worker_loop_cuda, daemon=True)
        else:
            self._worker_thread = Thread(target=self._worker_loop_single, daemon=True)
        # 工作队列：每个元素是 (message, from_node, upstream_event)
        self._work_queue = queue.Queue(maxsize=maxsize)
        # 启动工作线程
        self._worker_thread_init_event = Event()

        self._worker_thread.start()
        self._worker_thread_init_event.wait()
        self._stop_event = Event()

    def worker_loop_init(self):
        pass

    def stop(self):
        pass

    def message_pre_process(self, message: 'Message'):
        return message

    def message_post_process(self, message: 'Message'):
        return message

    def _worker_loop_cuda(self):
        try:
            torch.cuda.set_device(self.device)  # 线程局部设置设备
            self.worker_loop_init()
            stream = torch.cuda.Stream(device=self.device)
        finally:
            self._worker_thread_init_event.set()
        while True:
            # 从队列获取：消息 + 来源节点名 + 上游事件（可能为 None）
            message, from_node, upstream_event = self._work_queue.get()

            if message.type == 'stop':
                # 广播 stop 消息（无 GPU 依赖）
                for out_node in self.output_nodes:
                    out_node.input(message, from_node=self.name)
                self.stop()
                self._work_queue.task_done()
                self._stop_event.set()
                break
            try:
                # ✅ 在工作线程中等待上游完成（如果存在）
                if upstream_event is not None:
                    stream.wait_event(upstream_event)

                # 执行计算
                with torch.cuda.stream(stream):
                    message = self.message_pre_process(message)
                    self.process([message], stream=stream)
                    message = self.message_post_process(message)

                # 记录本节点完成事件
                local_event = torch.cuda.Event()
                local_event.record(stream)

                # 通知下游
                self._notify_downstream_cuda(message, local_event)

            except Exception as e:
                self.exception_callback(f'NODE_{self.name}', self, message, e)
                if self.logger:
                    self.logger.exception(f'{self.name} process error')
            finally:
                self._work_queue.task_done()

        torch.cuda.empty_cache()

    def _worker_loop_single(self):
        try:
            self.worker_loop_init()
        finally:
            self._worker_thread_init_event.set()

        while True:
            message, from_node, _ = self._work_queue.get()

            if message.type == 'stop':
                for out_node in self.output_nodes:
                    out_node.input(message, from_node=self.name)
                self.stop()
                self._work_queue.task_done()
                self._stop_event.set()
                break

            try:
                # 忽略 upstream_event（CPU 节点无需 GPU 同步）
                message = self.message_pre_process(message)
                self.process([message], stream=None)  # 注意：无 stream 参数
                message = self.message_post_process(message)

                # 通知下游（不传递事件）
                self._notify_downstream_single(message)

            except Exception as e:
                self.exception_callback(f'NODE_{self.name}', self, message, e)
                if self.logger:
                    self.logger.exception(f'{self.name} process error')
            finally:
                self._work_queue.task_done()

        # 可选：清理（CPU 场景通常不需要）
        pass

    def _notify_downstream_single(self, message: 'Message'):
        """纯 CPU 通知，不传递事件"""
        for out_node in self.output_nodes:
            if self.logger:
                self.logger.debug(
                    f'NODE_DATA [{self.name}-->{out_node.name}] '
                    f'{message.seq:06d} {message.timestamp}'
                )
            try:
                # 兼容 CUDA 节点：但 CPU 节点不传 upstream_event
                # 因此调用时只传两个参数（符合 input 签名的默认行为）
                out_node.input(message, from_node=self.name)
            except Exception as e:
                self.exception_callback(f'NODE_{self.name}', self, message, e)
                if self.logger:
                    self.logger.exception(f'{self.name} notify downstream error')

    def _notify_downstream_cuda(self, message: 'Message', local_event: torch.cuda.Event):
        """将消息和完成事件传递给所有下游节点"""
        for out_node in self.output_nodes:
            if self.logger:
                self.logger.debug(
                    f'NODE_DATA_STREAM [{self.name}-->{out_node.name}] '
                    f'{message.seq:06d} {message.timestamp}'
                )
            try:
                if isinstance(out_node, QueueCudaStreamNode):
                    # 下游是 CUDA 节点：传递事件，实现异步依赖
                    out_node.input(message, from_node=self.name, upstream_event=local_event)
                else:
                    # 下游是 CPU 节点：同步后传递
                    local_event.synchronize()
                    out_node.input(message, from_node=self.name)
            except Exception as e:
                self.exception_callback(f'NODE_{self.name}', self, message, e)
                if self.logger:
                    self.logger.exception(f'{self.name} notify downstream error')

    def input(self, message: 'Message', from_node: str, upstream_event: Optional[torch.cuda.Event] = None):
        """
        主线程入口：快速入队，不执行任何 CUDA 操作。

        参数:
            message: 输入消息
            from_node: 消息来源节点名
            upstream_event: （可选）上游完成事件，用于建立 GPU 依赖
        """

        self._work_queue.put((message, from_node, upstream_event))
        if message.type == 'stop':
            self._stop_event.wait(5)

    def process(self, messages: List['Message'], stream: torch.cuda.Stream):
        """
        子类必须实现此方法。
        在指定 CUDA 流中处理消息。
        """
        raise NotImplementedError("Subclasses must implement the 'process' method.")


class EventBus:
    def __init__(self):
        self._subscribers = {}  # topic -> [callbacks]

    def subscribe(self, topic, callback):
        if topic not in self._subscribers:
            self._subscribers[topic] = []
        self._subscribers[topic].append(callback)

    def publish(self, topic, *args, **kwargs):
        if topic in self._subscribers:
            for callback in self._subscribers[topic]:
                callback(*args, **kwargs)


class NodesChain:

    def __init__(self, context=None, logger=None, node_exception_callback=None):
        if context is None:
            self.context = {}
        self.message_maker = MessageMaker()
        self.ori_logger = logger
        self.logger_str = None
        if logger is None:
            logger = get_logger(__name__)
        elif isinstance(logger, str):
            self.logger_str = logger
            logger = get_logger(f"{logger}.chain")
        elif isinstance(logger, logging.Logger):
            logger = get_logger(f"{logger.name}.chain")
        else:
            logger = get_logger(__name__)
        self.logger = logger
        self.root_node = self.build()
        self.node_exception_callback = node_exception_callback
        self.event_bus = EventBus()
        self.__set_context__(self.root_node)

    def build(self) -> Node:
        pass

    def input(self, data):
        start_time = time.time()
        if isinstance(data, Message):
            self.root_node.input(data, 'root')
        else:
            self.root_node.input(self.message_maker.data_message(data), 'root')
        self.logger.debug(f'CHAIN once use time: {time.time() - start_time:.3f}s')

    def stop(self):
        self.root_node.input(self.message_maker.stop_message(), 'root')

    def __set_context__(self, node: Node):
        node.set_context(self.context)
        node.set_logger(self.ori_logger)
        node.event_bus = self.event_bus
        if self.node_exception_callback is not None:
            node.exception_callback = self.node_exception_callback
        for out in node.output_nodes:
            self.__set_context__(out)

    # def __draw_node__(self, node: Node, dot):
    #     dot.node(node.name, node.name)
    #     for out in node.output_nodes:
    #         self.__draw_node__(out, dot)
    #         dot.edge(node.name, out.name)
    # def graph_image(self, file_path):
    #     from graphviz import Digraph
    #     dot = Digraph(comment='chain')
    #     self.__draw_node__(self.root_node, dot)
    #     dot.render(file_path,format=os.path.splitext(file_path)[1][1:], cleanup=True)
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.stop()


def split_node(s: str):
    match = re.match(r'^(\w+)\[(\S+)]$', s.strip())
    if match:
        prefix, chinese = match.groups()
        return (prefix, chinese)
    else:
        # 没有括号时，第二项使用原字符串不变
        return (s, s)


class MermaidNodesChain(NodesChain):

    def __init__(self, context=None, logger=None, node_exception_callback=None):
        super().__init__(context=context, logger=logger, node_exception_callback=node_exception_callback)

    def build(self) -> Node:
        root_node, nodes_builder = self.nodes_builder()
        # nodes_map = dict({(node.name, node) for node in nodes})

        link_edges = self.link_edges()
        if isinstance(link_edges, str):
            lines = link_edges.strip().split('\n')
            edges = set()
            for line in lines:
                if "-->" not in line or line.strip().startswith('##') or line.strip().startswith('%%'):
                    continue
                nodes = line.strip().split('-->')
                nodes = [node.strip() for node in nodes if not node.strip().startswith("#")]
                # nodes = list(map(lambda x: re.sub(r'\[.*?]', '', x), nodes))
                nodes = list(map(split_node, nodes))
                edges.update([(k, v) for k, v in zip(nodes[:-1], nodes[1:])])
            link_edges = edges

        node_names = {}
        for from_node, to_node in link_edges:
            if isinstance(from_node, str):
                node_names[from_node.strip()] = from_node.strip()
            else:
                node_names[from_node[0]] = from_node[1]
            if isinstance(to_node, str):
                node_names[to_node.strip()] = to_node.strip()
            else:
                node_names[to_node[0]] = to_node[1]

        nodes_map = {}
        for node_name, node_description in node_names.items():
            nodes_map[node_name] = nodes_builder[node_name]()
            nodes_map[node_name].name = node_name
            nodes_map[node_name].description = node_description

        root_node = nodes_map[root_node]

        for from_node, to_node in link_edges:
            if isinstance(from_node, str):
                from_node = from_node.strip()
            else:
                from_node = from_node[0]
            if isinstance(to_node, str):
                to_node = to_node.strip()
            else:
                to_node = to_node[0]
            from_node, to_node = from_node.strip(), to_node.strip()
            from_node, to_node = nodes_map[from_node], nodes_map[to_node]
            from_node.registry_output(to_node)
        return root_node

    def nodes_builder(self):  # -> (str, map)
        pass

    def link_edges(self):  # -> set[tuple] | str
        pass


class NumberAdd(SingleInputNode):

    def __init__(self, context=None, name=None, num=0):
        self.num = num
        super().__init__(context=context, name=name)

    def process(self, message: Message, **kwargs):
        message.payload['data'] = message.payload['data'] + self.num


class PrintNode(SingleInputNode):

    def __init__(self, context=None, name='print'):
        super().__init__(context=context, name=name)

    def process(self, message: Message, **kwargs):
        print(message.payload)


class TestChain(NodesChain):

    def build(self) -> Node:
        add1 = NumberAdd(name='add_1', num=1)
        add5 = NumberAdd(name='add_5', num=5)
        print_str = PrintNode(name='print')

        add1.registry_output(add5)
        add5.registry_output(print_str)

        return add1


if __name__ == '__main__':
    from aabd.base.log_setting import set_logger

    set_logger('chain_nodes', log_level='DEBUG')
    chain = TestChain()
    # chain.graph_image('1.png')
    chain.input({'data': 1})
    chain.stop()
