import abc
import time
import asyncio
from pydantic import BaseModel
from typing import Optional, List, Dict, Any
from aworld.core.event.base import Message
from enum import Enum
from abc import ABC, abstractmethod, ABCMeta
from aworld.core.agent.base import is_agent_by_name
from aworld.core.tool.tool_desc import is_tool_by_name
from aworld.core.singleton import InheritanceSingleton, SingletonMeta
from aworld.core.event.base import Constants
from aworld.logs.util import logger
from aworld.events.util import send_message
from aworld.utils.serialized_util import to_serializable


class RunNodeBusiType(Enum):
    AGENT = 'AGENT'
    TOOL = 'TOOL'
    TASK = 'TASK'
    TOOL_CALLBACK = 'TOOL_CALLBACK'
    HUMAN = 'HUMAN'
    CONTEXT = "CONTEXT"
    MEMORY = 'MEMORY'

    @staticmethod
    def from_message_category(category: str) -> 'RunNodeBusiType':
        if category == Constants.AGENT:
            return RunNodeBusiType.AGENT
        if category == Constants.TOOL:
            return RunNodeBusiType.TOOL
        if category == Constants.TASK:
            return RunNodeBusiType.TASK
        if category == Constants.TOOL_CALLBACK:
            return RunNodeBusiType.TOOL_CALLBACK
        if category == Constants.HUMAN:
            return RunNodeBusiType.HUMAN
        if category == Constants.MEMORY:
            return RunNodeBusiType.MEMORY
        if category == Constants.CONTEXT:
            return RunNodeBusiType.CONTEXT
        return None


class RunNodeStatus(Enum):
    INIT = 'INIT'
    RUNNING = 'RUNNING'
    BREAKED = 'BREAKED'
    SUCCESS = 'SUCCESS'
    FAILED = 'FAILED'
    TIMEOUT = 'TIMEOUT'


class HandleResult(BaseModel):
    name: str = None
    status: RunNodeStatus = None
    result_msg: Optional[str] = None
    result: Optional[Message] = None

    def to_dict(self) -> Dict[str, Any]:
        return {
            "name": self.name,
            "status": self.status.value,
            "result_msg": self.result_msg,
            "result": to_serializable(self.result) if self.result else None
        }


class RunNode(BaseModel):
    # {busi_id}_{busi_type}
    node_id: Optional[str] = None
    task_id: Optional[str] = None
    busi_type: str = None
    busi_id: Optional[str] = None
    session_id: Optional[str] = None
    msg_id: Optional[str] = None  # input message id
    # busi_id of node that send the input message
    msg_from: Optional[str] = None
    parent_node_id: Optional[str] = None
    status: RunNodeStatus = None
    result_msg: Optional[str] = None
    results: Optional[List[HandleResult]] = None
    create_time: Optional[float] = None
    execute_time: Optional[float] = None
    end_time: Optional[float] = None
    group_id: Optional[str] = None
    # sub_group_root_id required when group_id is not None
    sub_group_root_id: Optional[str] = None
    # metadata is used to store the context of the sub task when group_id is not None
    metadata: Optional[dict] = None

    def has_finished(self):
        return self.status in [RunNodeStatus.SUCCESS, RunNodeStatus.FAILED, RunNodeStatus.TIMEOUT]


class SubGroup(BaseModel):
    '''
    SubGroup represents an execution chain pointing to the root node 
    '''
    root_node_id: Optional[str] = None
    session_id: Optional[str] = None
    group_id: Optional[str] = None
    create_time: Optional[float] = None
    execute_time: Optional[float] = None
    end_time: Optional[float] = None
    status: RunNodeStatus = None
    result_msg: Optional[str] = None
    results: Optional[List[HandleResult]] = None
    metadata: Optional[dict] = None

    def has_finished(self):
        return self.status in [RunNodeStatus.SUCCESS, RunNodeStatus.FAILED, RunNodeStatus.TIMEOUT]


class NodeGroup(BaseModel):
    '''
    Node group, used to manage sub group
    '''
    group_id: str = None
    session_id: str = None
    # subtask root node id list
    root_node_ids: List[str] = None
    finished: Optional[bool] = False
    finish_notified: Optional[bool] = False
    create_time: Optional[float] = None
    execute_time: Optional[float] = None
    end_time: Optional[float] = None
    status: RunNodeStatus = None
    # failed subtask root node id list
    failed_root_node_ids: Optional[List[str]] = None
    parent_group_id: Optional[str] = None
    metadata: Optional[dict] = None

    def has_finished(self):
        return self.status in [RunNodeStatus.SUCCESS, RunNodeStatus.FAILED, RunNodeStatus.TIMEOUT]


class NodeGroupDetail(NodeGroup):
    sub_groups: Optional[List[SubGroup]] = None


class StateStorage:
    __metaclass__ = abc.ABCMeta

    @abstractmethod
    def get(self, node_id: str) -> RunNode:
        pass

    @abstractmethod
    def insert(self, node: RunNode):
        pass

    @abstractmethod
    def update(self, node: RunNode):
        pass

    @abstractmethod
    def query(self, session_id: str) -> List[RunNode]:
        pass

    @abstractmethod
    def query_by_task_id(self, task_id: str) -> List[RunNode]:
        pass


class NodeGroupStorage:
    __metaclass__ = abc.ABCMeta

    @abstractmethod
    def get(self, group_id: str) -> NodeGroup:
        pass

    @abstractmethod
    def insert(self, node_group: NodeGroup):
        pass

    @abstractmethod
    def update(self, node_group: NodeGroup):
        pass


class SubGroupStorage:
    __metaclass__ = abc.ABCMeta

    @abstractmethod
    def get(self, node_id: str) -> SubGroup:
        pass

    @abstractmethod
    def insert(self, sub_group: SubGroup):
        pass

    @abstractmethod
    def update(self, sub_group: SubGroup):
        pass


class StateStorageMeta(SingletonMeta, ABCMeta):
    pass


class InMemoryStateStorage(StateStorage, InheritanceSingleton, metaclass=StateStorageMeta):
    '''
    In memory state storage
    '''

    def __init__(self, max_session=1000):
        self._max_session = max_session
        self._nodes = {}  # {node_id: RunNode}
        self._ordered_session_ids = []
        self._session_nodes = {}  # {session_id: [RunNode, RunNode]}

    def get(self, node_id: str) -> RunNode:
        return self._nodes.get(node_id)

    def insert(self, node: RunNode):
        if node.session_id not in self._ordered_session_ids:
            self._ordered_session_ids.append(node.session_id)
            self._session_nodes.update({node.session_id: []})
        if node.node_id not in self._nodes:
            self._nodes.update({node.node_id: node})
            self._session_nodes[node.session_id].append(node)

        if len(self._ordered_session_ids) > self._max_session:
            oldest_session_id = self._ordered_session_ids.pop(0)
            session_nodes = self._session_nodes.pop(oldest_session_id)
            for node in session_nodes:
                self._nodes.pop(node.node_id)
        # logger.info(f"storage nodes: {self._nodes}")

    def update(self, node: RunNode):
        self._nodes[node.node_id] = node

    def query(self, session_id: str, msg_id: str = None) -> List[RunNode]:
        session_nodes = self._session_nodes.get(session_id, [])
        if msg_id:
            return [node for node in session_nodes if node.msg_id == msg_id]
        return session_nodes

    def query_by_task_id(self, task_id: str) -> List[RunNode]:
        return [node for node in self._nodes.values() if node.task_id == task_id]


class InMemoryNodeGroupStorage(NodeGroupStorage, InheritanceSingleton, metaclass=StateStorageMeta):
    '''
    In memory node group storage
    '''

    def __init__(self):
        self.node_groups = {}

    def get(self, group_id: str) -> NodeGroup:
        return self.node_groups.get(group_id)

    def insert(self, node_group: NodeGroup):
        self.node_groups[node_group.group_id] = node_group

    def update(self, node_group: NodeGroup):
        self.node_groups[node_group.group_id] = node_group


class InMemorySubGroupStorage(SubGroupStorage, InheritanceSingleton, metaclass=StateStorageMeta):
    '''
    In memory sub task storage
    '''

    def __init__(self):
        self.sub_groups = {}

    def get(self, node_id: str) -> SubGroup:
        return self.sub_groups.get(node_id)

    def insert(self, sub_group: SubGroup):
        self.sub_groups[sub_group.root_node_id] = sub_group

    def update(self, sub_group: SubGroup):
        self.sub_groups[sub_group.root_node_id] = sub_group


class RuntimeStateManager(InheritanceSingleton):
    '''
    Runtime state manager
    '''

    def __init__(self,
                 storage: StateStorage = InMemoryStateStorage.instance()):
        self.storage = storage
        self._node_group_manager = None

    @property
    def node_group_manager(self):
        if not self._node_group_manager:
            self._node_group_manager = NodeGroupManager(node_state_manager=self)
        return self._node_group_manager

    def create_node(self,
                    busi_type: RunNodeBusiType,
                    busi_id: str,
                    session_id: str,
                    node_id: str = None,
                    task_id: str = None,
                    parent_node_id: str = None,
                    msg_id: str = None,
                    msg_from: str = None,
                    group_id: str = None,
                    sub_group_root_id: str = None,
                    metadata: Optional[dict] = None) -> RunNode:
        '''
            create node and insert to storage
        '''
        node_id = node_id or msg_id
        node = self._find_node(node_id)
        if node:
            # raise Exception(f"node already exist, node_id: {node_id}")
            return
        if parent_node_id:
            parent_node = self._find_node(parent_node_id)
            if not parent_node:
                logger.warning(
                    f"parent node not exist, parent_node_id: {parent_node_id}")
        node = RunNode(node_id=node_id,
                       busi_type=busi_type.name,
                       busi_id=busi_id,
                       session_id=session_id,
                       task_id=task_id,
                       msg_id=msg_id,
                       msg_from=msg_from,
                       parent_node_id=parent_node_id,
                       status=RunNodeStatus.INIT,
                       create_time=time.time(),
                       group_id=group_id,
                       sub_group_root_id=sub_group_root_id,
                       metadata=metadata)
        self.storage.insert(node)
        # create sub group if node is the root node of sub group
        if group_id and sub_group_root_id and node_id == sub_group_root_id:
            self.node_group_manager.create_sub_group(group_id, session_id, sub_group_root_id, metadata)
        return node

    def run_node(self, node_id: str):
        '''
            set node status to RUNNING and update to storage
        '''
        logger.debug(f"set node {node_id} running")
        node = self._node_exist(node_id)
        node.status = RunNodeStatus.RUNNING
        node.execute_time = time.time()
        self.storage.update(node)
        # update sub group status if node is the root node
        if node.group_id and node.sub_group_root_id and node.node_id == node.sub_group_root_id:
            self.node_group_manager.run_sub_group(node_id)

    def save_result(self,
                    node_id: str,
                    result: HandleResult):
        '''
            save node execute result and update to storage
        '''
        node = self._node_exist(node_id)
        if not node.results:
            node.results = []
        node.results.append(result)
        self.storage.update(node)

    def break_node(self, node_id):
        '''
            set node status to BREAKED and update to storage
        '''
        node = self._node_exist(node_id)
        node.status = RunNodeStatus.BREAKED
        self.storage.update(node)

    def run_succeed(self,
                    node_id,
                    result_msg=None,
                    results: List[HandleResult] = None):
        '''
            set node status to SUCCESS and update to storage
        '''
        node = self._node_exist(node_id)
        node.status = RunNodeStatus.SUCCESS
        node.result_msg = result_msg
        node.end_time = time.time()
        if results:
            if not node.results:
                node.results = []
            node.results.extend(results)
        logger.debug(f"set node {node_id} succeed: {node}")

        self.storage.update(node)

    def run_failed(self,
                   node_id,
                   result_msg=None,
                   results: List[HandleResult] = None):
        '''
            set node status to FAILED and update to storage
        '''
        node = self._node_exist(node_id)
        node.status = RunNodeStatus.FAILED
        node.result_msg = result_msg
        node.end_time = time.time()
        if results:
            if not node.results:
                node.results = []
            node.results.extend(results)
        self.storage.update(node)

    def run_timeout(self,
                    node_id,
                    result_msg=None):
        '''
            set node status to TIMEOUT and update to storage
        '''
        node = self._node_exist(node_id)
        node.status = RunNodeStatus.TIMEOUT
        node.result_msg = result_msg
        self.storage.update(node)

    def finish_sub_task(self, node_id: str):
        '''
            finish sub task with node_id as the root node
        '''
        node = self._node_exist(node_id)
        node.sub_task_finished = True
        self.storage.update(node)

    def get_node(self, node_id: str) -> RunNode:
        '''
            get node from storage
        '''
        return self._find_node(node_id)

    def get_nodes(self, session_id: str) -> List[RunNode]:
        '''
            get nodes from storage
        '''
        return self.storage.query(session_id)

    def _node_exist(self, node_id: str):
        node = self._find_node(node_id)
        if not node:
            raise Exception(f"node not found, node_id: {node_id}")
        return node

    def _find_node(self, node_id: str):
        return self.storage.get(node_id)

    def _judge_msg_from_busi_type(self, msg_from: str) -> RunNodeBusiType:
        '''
        judge msg_from busi_type
        '''
        if is_agent_by_name(msg_from):
            return RunNodeBusiType.AGENT
        if is_tool_by_name(msg_from):
            return RunNodeBusiType.TOOL
        return RunNodeBusiType.TASK

    async def wait_for_node_completion(self, node_id: str, timeout: float = 120.0, interval: float = 1.0) -> RunNode:
        '''Poll for node status until completion or timeout.

        Args:
            node_id: Node ID
            timeout: Timeout threshold in seconds
            interval: Polling interval in seconds

        Returns:
            RunNode: Node object

        Raises:
            Exception: If node does not exist
            TimeoutError: If waiting times out
        '''
        start_time = time.time()
        log_start_time = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
        logger.info(f"wait for node completion: {node_id}, start_time:{log_start_time}")

        while True:
            node = self._find_node(node_id)
            if not node:
                raise Exception(f"Node not found, node_id: {node_id}")

            # Check if node has completed
            if node.status in [RunNodeStatus.SUCCESS, RunNodeStatus.FAILED, RunNodeStatus.BREAKED,
                               RunNodeStatus.TIMEOUT]:
                return node

            # Check if timed out
            if time.time() - start_time > timeout:
                self.run_timeout(node_id, result_msg=f"Waiting for node completion timed out after {timeout} seconds")
                node = self._find_node(node_id)
                logger.warn(f"wait for node completion timed out: {node_id}, node: {node}")
                return node

            # Wait for the specified interval before polling again
            await asyncio.sleep(interval)

    async def create_group(self, group_id: str,
                           session_id: str,
                           root_node_ids: List[str] = None,
                           parent_group_id: Optional[str] = None,
                           metadata: Optional[dict] = None) -> NodeGroup:
        '''
        create node group
        '''
        return await self.node_group_manager.create_group(group_id, session_id, root_node_ids, parent_group_id, metadata)

    async def finish_sub_group(self,
                               group_id: str,
                               root_node_id: str,
                               results: List[Message] = None,
                               result_msg: str = None):
        '''
        finish sub group
        '''
        handle_results = []
        for msg in results:
            handle_result = HandleResult(
                status=RunNodeStatus.FAILED if msg.is_error() else RunNodeStatus.SUCCESS,
                result=msg,
                name=msg.sender
            )
            handle_results.append(handle_result)
        await self.node_group_manager.finish_sub_group(group_id, root_node_id, handle_results, result_msg)

    def get_group(self, group_id: str) -> NodeGroup:
        '''
        get group basic info
        '''
        return self.node_group_manager.get_group(group_id)

    def query_group_detail(self, group_id: str) -> NodeGroupDetail:
        '''
        query group detail info with all sub group info
        '''
        return self.node_group_manager.query_group_detail(group_id)

    def query_by_task(self, task_id: str, busi_typ: RunNodeBusiType = None, busi_id: str = None) -> List[RunNode]:
        all_task_nodes = self.storage.query_by_task_id(task_id)
        if (not busi_typ and busi_id) or (busi_typ and not busi_id):
            raise Exception("busi_typ and busi_id must be both None or not None")
        if busi_typ and busi_id:
            result_nodes = [node for node in all_task_nodes if node.busi_type == busi_typ.name and node.busi_id == busi_id]
        else:
            result_nodes = all_task_nodes
        result_nodes.sort(key=lambda x: x.create_time if x.create_time else 0, reverse=True)
        return result_nodes


class NodeGroupManager(InheritanceSingleton):
    '''
    Node group manager, used to manage node group
    '''

    def __init__(self,
                 sub_group_storage: SubGroupStorage = InMemorySubGroupStorage.instance(),
                 node_group_storage: NodeGroupStorage = InMemoryNodeGroupStorage.instance(),
                 node_state_manager: RuntimeStateManager = None):
        self.sub_group_storage = sub_group_storage
        self.node_group_storage = node_group_storage
        self.node_state_manager = node_state_manager

    async def create_group(self, group_id: str,
                           session_id: str,
                           root_node_ids: List[str] = None,
                           parent_group_id: Optional[str] = None,
                           metadata: Optional[dict] = None) -> NodeGroup:
        '''
        create node group
        '''
        group = self._find_group(group_id)
        if group:
            raise Exception(f"group already exist, group_id: {group_id}")
        node_group = NodeGroup(
            session_id=session_id,
            group_id=group_id,
            root_node_ids=root_node_ids,
            parent_group_id=parent_group_id,
            metadata=metadata,
            create_time=time.time(),
            update_time=time.time(),
            status=RunNodeStatus.INIT,
        )
        self.node_group_storage.insert(node_group)
        await self._check_subgroup_status(group_id, root_node_ids)

    def create_sub_group(self,
                         group_id: str,
                         session_id: str,
                         root_node_id: str,
                         metadata: Optional[dict] = None) -> SubGroup:
        '''
        create sub group
        '''
        subgroup = self._find_subgroup(root_node_id)
        if subgroup:
            raise Exception(f"subgroup already exist, group_id: {group_id}, root_node_id: {root_node_id}")
        run_node = self.node_state_manager.get_node(root_node_id)
        if not run_node:
            raise Exception(f"run node not found, root_node_id: {root_node_id}")

        sub_group = SubGroup(
            session_id=session_id,
            group_id=group_id,
            root_node_id=root_node_id,
            metadata=metadata,
            create_time=time.time(),
            update_time=time.time(),
            status=RunNodeStatus.INIT,
        )
        self.sub_group_storage.insert(sub_group)
        return sub_group

    def run_sub_group(self,
                      root_node_id: str):
        '''
        run sub group
        '''
        subgroup = self._subgroup_exist(root_node_id)
        if not subgroup:
            raise Exception(f"subgroup not found, root_node_id: {root_node_id}")

        subgroup.execute_time = time.time()
        subgroup.status = RunNodeStatus.RUNNING
        self.sub_group_storage.update(subgroup)
        self.run_group(subgroup.group_id)

    def run_group(self, group_id):
        group = self.node_group_storage.get(group_id)
        if group.status == RunNodeStatus.INIT:
            group.status = RunNodeStatus.RUNNING
            group.execute_time = time.time()
            self.node_group_storage.update(group)

    async def finish_sub_group(self,
                               group_id: str,
                               root_node_id: str,
                               results: List[HandleResult] = None,
                               result_msg: str = None):
        '''
        finish sub task with node_id as the root node
        '''
        subgroup = self.sub_group_storage.get(root_node_id)
        if not subgroup:
            raise Exception(f"subgroup not found, group_id: {group_id}, root_node_id: {root_node_id}")
        if subgroup.group_id != group_id:
            raise Exception(f"subgroup group_id not match, group_id: {group_id}, root_node_id: {root_node_id}")

        group = self._group_exist(group_id)
        subgroup.end_time = time.time()
        subgroup.results = results
        subgroup.result_msg = result_msg
        subgroup.status = RunNodeStatus.SUCCESS
        for result in results:
            if result.status == RunNodeStatus.FAILED:
                subgroup.status = RunNodeStatus.FAILED
        self.sub_group_storage.update(subgroup)
        # check all subgroup status and update group status
        await self._check_subgroup_status(group_id, group.root_node_ids)

    async def _check_subgroup_status(self, group_id, root_node_ids: List[str]):
        '''
        check subgroups status and update group status, if group finished, send group finish message
        '''
        all_subgroups_finished = True
        failed_subgroups = []
        for root_node_id in root_node_ids:
            subgroup = self.sub_group_storage.get(root_node_id)
            if not subgroup or not subgroup.has_finished():
                all_subgroups_finished = False
                break
            if subgroup.status == RunNodeStatus.FAILED or subgroup.status == RunNodeStatus.TIMEOUT:
                failed_subgroups.append(subgroup)

        if all_subgroups_finished:
            group = self._group_exist(group_id)
            if failed_subgroups:
                group.status = RunNodeStatus.FAILED
                group.failed_root_node_ids = [subgroup.root_node_id for subgroup in failed_subgroups]
            else:
                group.status = RunNodeStatus.SUCCESS
            group.end_time = time.time()
            self.node_group_storage.update(group)
            await self._send_group_finish_message(group_id)

    async def _send_group_finish_message(self, group_id: str):
        '''
            Currently, for simple implementation, concurrency control needs to be considered in a distributed environment
        '''
        group = self._group_exist(group_id)
        if group.finish_notified:
            logger.warning(f"group finish message already sent, group_id: {group_id}")
            return
        group_results = {}
        metadata = group.metadata
        for root_node_id in group.root_node_ids:
            subgroup = self.sub_group_storage.get(root_node_id)
            group_results[root_node_id] = subgroup.results
            if not metadata:
                metadata = subgroup.metadata

        metadata = metadata or {}
        if group.parent_group_id:
            metadata.update({
                "parent_group_id": group.parent_group_id
            })
        message = Message(
            category="group",
            payload=group_results,
            sender="node_group_manager",
            session_id=group.session_id,
            topic="__group_results",
            headers=metadata
        )
        await send_message(message)
        group.finish_notified = True
        self.node_group_storage.update(group)

    def get_group(self, group_id: str) -> NodeGroup:
        '''
            get group basic info
        '''
        return self._find_group(group_id)

    def query_group_detail(self, group_id: str) -> NodeGroupDetail:
        '''
            query group detail info with all sub group info
        '''
        group = self._find_group(group_id)
        if not group:
            return None
        sub_groups = []
        for root_node_id in group.root_node_ids:
            subgroup = self._find_subgroup(root_node_id)
            if subgroup:
                sub_groups.append(subgroup)
        return NodeGroupDetail(
            group_id=group.group_id,
            root_node_ids=group.root_node_ids,
            parent_group_id=group.parent_group_id,
            metadata=group.metadata,
            create_time=group.create_time,
            execute_time=group.execute_time,
            end_time=group.end_time,
            status=group.status,
            failed_root_node_ids=group.failed_root_node_ids,
            sub_groups=sub_groups
        )

    def _find_subgroup(self, root_node_id: str) -> SubGroup:
        return self.sub_group_storage.get(root_node_id)

    def _subgroup_exist(self, root_node_id: str) -> SubGroup:
        subgroup = self._find_subgroup(root_node_id)
        if not subgroup:
            raise Exception(f"subgroup not found, root_node_id: {root_node_id}")
        return subgroup

    def _find_group(self, group_id: str) -> NodeGroup:
        return self.node_group_storage.get(group_id)

    def _group_exist(self, group_id: str) -> NodeGroup:
        group = self._find_group(group_id)
        if not group:
            raise Exception(f"group not found, group_id: {group_id}")
        return group


class EventRuntimeStateManager(RuntimeStateManager):

    def __init__(self, storage: StateStorage = InMemoryStateStorage.instance()):
        super().__init__(storage)

    def start_message_node(self, message: Message):
        '''
        create and start node while message handle started.
        '''
        metadata = message.headers
        run_node_busi_type = RunNodeBusiType.from_message_category(
            message.category)
        logger.debug(
            f"start message node: {message.receiver}, busi_type={run_node_busi_type}, node_id={message.id}")
        if run_node_busi_type:
            self.create_node(
                node_id=message.id,
                busi_type=run_node_busi_type,
                busi_id=message.receiver or "",
                session_id=message.session_id,
                task_id=message.task_id,
                msg_id=message.id,
                msg_from=message.sender,
                group_id=metadata.get("group_id") if metadata else None,
                sub_group_root_id=metadata.get("root_message_id") if metadata else None,
                metadata=metadata)
            self.run_node(message.id)

    def save_message_handle_result(self, name: str, message: Message, result: Message = None):
        '''
        save message handle result
        '''
        run_node_busi_type = RunNodeBusiType.from_message_category(
            message.category)
        if run_node_busi_type:
            if result and result.is_error():
                handle_result = HandleResult(
                    name=name,
                    status=RunNodeStatus.FAILED,
                    result=result)
            else:
                handle_result = HandleResult(
                    name=name,
                    status=self.get_node(message.id).status if self.get_node(message.id) else RunNodeStatus.FAILED,
                    result=result)
            self.save_result(node_id=message.id, result=handle_result)

    def end_message_node(self, message: Message):
        '''
        end node while message handle finished.
        '''
        run_node_busi_type = RunNodeBusiType.from_message_category(
            message.category)
        if run_node_busi_type:
            node = self._node_exist(node_id=message.id)
            status = RunNodeStatus.SUCCESS
            if node.results:
                for result in node.results:
                    if result.status == RunNodeStatus.FAILED:
                        status = RunNodeStatus.FAILED
                        break
            if status == RunNodeStatus.FAILED:
                self.run_failed(node_id=message.id)
            else:
                self.run_succeed(node_id=message.id)

    def get_message_node_status(self, message: Message) -> RunNodeStatus:
        node = self.get_node(node_id=message.id)
        if not node:
            return RunNodeStatus.INIT
        return node.status
