import logging
import queue
import time
import uuid
from collections.abc import Generator, Mapping
from concurrent.futures import ThreadPoolExecutor, wait
from copy import copy, deepcopy
from datetime import UTC, datetime
from typing import Any, Optional, cast

from flask import Flask, current_app

from configs import dify_config
from core.app.apps.base_app_queue_manager import GenerateTaskStoppedError
from core.app.entities.app_invoke_entities import InvokeFrom
from core.workflow.entities.node_entities import NodeRunMetadataKey, NodeRunResult
from core.workflow.entities.variable_pool import VariablePool, VariableValue
from core.workflow.graph_engine.condition_handlers.condition_manager import ConditionManager
from core.workflow.graph_engine.entities.event import (
    BaseIterationEvent,
    GraphEngineEvent,
    GraphRunFailedEvent,
    GraphRunPartialSucceededEvent,
    GraphRunStartedEvent,
    GraphRunSucceededEvent,
    NodeRunExceptionEvent,
    NodeRunFailedEvent,
    NodeRunRetrieverResourceEvent,
    NodeRunRetryEvent,
    NodeRunStartedEvent,
    NodeRunStreamChunkEvent,
    NodeRunSucceededEvent,
    ParallelBranchRunFailedEvent,
    ParallelBranchRunStartedEvent,
    ParallelBranchRunSucceededEvent,
)
from core.workflow.graph_engine.entities.graph import Graph, GraphEdge
from core.workflow.graph_engine.entities.graph_init_params import GraphInitParams
from core.workflow.graph_engine.entities.graph_runtime_state import GraphRuntimeState
from core.workflow.graph_engine.entities.runtime_route_state import RouteNodeState
from core.workflow.nodes import NodeType
from core.workflow.nodes.answer.answer_stream_processor import AnswerStreamProcessor
from core.workflow.nodes.answer.base_stream_processor import StreamProcessor
from core.workflow.nodes.base import BaseNode
from core.workflow.nodes.base.entities import BaseNodeData
from core.workflow.nodes.end.end_stream_processor import EndStreamProcessor
from core.workflow.nodes.enums import ErrorStrategy, FailBranchSourceHandle
from core.workflow.nodes.event import RunCompletedEvent, RunRetrieverResourceEvent, RunStreamChunkEvent
from core.workflow.nodes.node_mapping import NODE_TYPE_CLASSES_MAPPING
from extensions.ext_database import db
from models.enums import UserFrom
from models.workflow import WorkflowNodeExecutionStatus, WorkflowType

logger = logging.getLogger(__name__)


class GraphEngineThreadPool(ThreadPoolExecutor):
    def __init__(
        self,
        max_workers=None,
        thread_name_prefix="",
        initializer=None,
        initargs=(),
        max_submit_count=dify_config.MAX_SUBMIT_COUNT,
    ) -> None:
        super().__init__(max_workers, thread_name_prefix, initializer, initargs)
        self.max_submit_count = max_submit_count
        self.submit_count = 0

    def submit(self, fn, /, *args, **kwargs):
        self.submit_count += 1
        self.check_is_full()

        return super().submit(fn, *args, **kwargs)

    def task_done_callback(self, future):
        self.submit_count -= 1

    def check_is_full(self) -> None:
        if self.submit_count > self.max_submit_count:
            raise ValueError(f"Max submit count {self.max_submit_count} of workflow thread pool reached.")


class GraphEngine:
    workflow_thread_pool_mapping: dict[str, GraphEngineThreadPool] = {}

    def __init__(
        self,
        tenant_id: str,
        app_id: str,
        workflow_type: WorkflowType,
        workflow_id: str,
        user_id: str,
        user_from: UserFrom,
        invoke_from: InvokeFrom,
        call_depth: int,
        graph: Graph,
        graph_config: Mapping[str, Any],
        variable_pool: VariablePool,
        max_execution_steps: int,
        max_execution_time: int,
        thread_pool_id: Optional[str] = None,
    ) -> None:
        thread_pool_max_submit_count = dify_config.MAX_SUBMIT_COUNT
        thread_pool_max_workers = 10

        # init thread pool
        if thread_pool_id:
            if thread_pool_id not in GraphEngine.workflow_thread_pool_mapping:
                raise ValueError(f"Max submit count {thread_pool_max_submit_count} of workflow thread pool reached.")

            self.thread_pool_id = thread_pool_id
            self.thread_pool = GraphEngine.workflow_thread_pool_mapping[thread_pool_id]
            self.is_main_thread_pool = False
        else:
            self.thread_pool = GraphEngineThreadPool(
                max_workers=thread_pool_max_workers, max_submit_count=thread_pool_max_submit_count
            )
            self.thread_pool_id = str(uuid.uuid4())
            self.is_main_thread_pool = True
            GraphEngine.workflow_thread_pool_mapping[self.thread_pool_id] = self.thread_pool

        self.graph = graph
        self.init_params = GraphInitParams(
            tenant_id=tenant_id,
            app_id=app_id,
            workflow_type=workflow_type,
            workflow_id=workflow_id,
            graph_config=graph_config,
            user_id=user_id,
            user_from=user_from,
            invoke_from=invoke_from,
            call_depth=call_depth,
        )

        self.graph_runtime_state = GraphRuntimeState(variable_pool=variable_pool, start_at=time.perf_counter())

        self.max_execution_steps = max_execution_steps
        self.max_execution_time = max_execution_time

    def run(self) -> Generator[GraphEngineEvent, None, None]:
        # trigger graph run start event
        yield GraphRunStartedEvent()
        handle_exceptions: list[str] = []
        stream_processor: StreamProcessor

        try:
            if self.init_params.workflow_type == WorkflowType.CHAT:
                stream_processor = AnswerStreamProcessor(
                    graph=self.graph, variable_pool=self.graph_runtime_state.variable_pool
                )
            else:
                stream_processor = EndStreamProcessor(
                    graph=self.graph, variable_pool=self.graph_runtime_state.variable_pool
                )

            # run graph
            generator = stream_processor.process(
                self._run(start_node_id=self.graph.root_node_id, handle_exceptions=handle_exceptions)
            )
            for item in generator:
                try:
                    yield item
                    if isinstance(item, NodeRunFailedEvent):
                        yield GraphRunFailedEvent(
                            error=item.route_node_state.failed_reason or "Unknown error.",
                            exceptions_count=len(handle_exceptions),
                        )
                        return
                    elif isinstance(item, NodeRunSucceededEvent):
                        if item.node_type == NodeType.END:
                            self.graph_runtime_state.outputs = (
                                dict(item.route_node_state.node_run_result.outputs)
                                if item.route_node_state.node_run_result
                                and item.route_node_state.node_run_result.outputs
                                else {}
                            )
                        elif item.node_type == NodeType.ANSWER:
                            if "answer" not in self.graph_runtime_state.outputs:
                                self.graph_runtime_state.outputs["answer"] = ""

                            self.graph_runtime_state.outputs["answer"] += "\n" + (
                                item.route_node_state.node_run_result.outputs.get("answer", "")
                                if item.route_node_state.node_run_result
                                and item.route_node_state.node_run_result.outputs
                                else ""
                            )

                            self.graph_runtime_state.outputs["answer"] = self.graph_runtime_state.outputs[
                                "answer"
                            ].strip()
                except Exception as e:
                    logger.exception("Graph run failed")
                    yield GraphRunFailedEvent(error=str(e), exceptions_count=len(handle_exceptions))
                    return
            # count exceptions to determine partial success
            if len(handle_exceptions) > 0:
                yield GraphRunPartialSucceededEvent(
                    exceptions_count=len(handle_exceptions), outputs=self.graph_runtime_state.outputs
                )
            else:
                # trigger graph run success event
                yield GraphRunSucceededEvent(outputs=self.graph_runtime_state.outputs)
            self._release_thread()
        except GraphRunFailedError as e:
            yield GraphRunFailedEvent(error=e.error, exceptions_count=len(handle_exceptions))
            self._release_thread()
            return
        except Exception as e:
            logger.exception("Unknown Error when graph running")
            yield GraphRunFailedEvent(error=str(e), exceptions_count=len(handle_exceptions))
            self._release_thread()
            raise e

    def _release_thread(self):
        if self.is_main_thread_pool and self.thread_pool_id in GraphEngine.workflow_thread_pool_mapping:
            del GraphEngine.workflow_thread_pool_mapping[self.thread_pool_id]

    def _run(
        self,
        start_node_id: str,
        in_parallel_id: Optional[str] = None,
        parent_parallel_id: Optional[str] = None,
        parent_parallel_start_node_id: Optional[str] = None,
        handle_exceptions: list[str] = [],
    ) -> Generator[GraphEngineEvent, None, None]:
        parallel_start_node_id = None
        if in_parallel_id:
            parallel_start_node_id = start_node_id

        next_node_id = start_node_id
        previous_route_node_state: Optional[RouteNodeState] = None
        while True:
            # max steps reached
            if self.graph_runtime_state.node_run_steps > self.max_execution_steps:
                raise GraphRunFailedError("Max steps {} reached.".format(self.max_execution_steps))

            # or max execution time reached
            if self._is_timed_out(
                start_at=self.graph_runtime_state.start_at, max_execution_time=self.max_execution_time
            ):
                raise GraphRunFailedError("Max execution time {}s reached.".format(self.max_execution_time))

            # init route node state
            route_node_state = self.graph_runtime_state.node_run_state.create_node_state(node_id=next_node_id)

            # get node config
            node_id = route_node_state.node_id
            node_config = self.graph.node_id_config_mapping.get(node_id)
            if not node_config:
                raise GraphRunFailedError(f"Node {node_id} config not found.")

            # convert to specific node
            node_type = NodeType(node_config.get("data", {}).get("type"))
            node_version = node_config.get("data", {}).get("version", "1")
            node_cls = NODE_TYPE_CLASSES_MAPPING[node_type][node_version]

            previous_node_id = previous_route_node_state.node_id if previous_route_node_state else None

            # init workflow run state
            node_instance = node_cls(  # type: ignore
                id=route_node_state.id,
                config=node_config,
                graph_init_params=self.init_params,
                graph=self.graph,
                graph_runtime_state=self.graph_runtime_state,
                previous_node_id=previous_node_id,
                thread_pool_id=self.thread_pool_id,
            )
            node_instance = cast(BaseNode[BaseNodeData], node_instance)
            try:
                # run node
                generator = self._run_node(
                    node_instance=node_instance,
                    route_node_state=route_node_state,
                    parallel_id=in_parallel_id,
                    parallel_start_node_id=parallel_start_node_id,
                    parent_parallel_id=parent_parallel_id,
                    parent_parallel_start_node_id=parent_parallel_start_node_id,
                    handle_exceptions=handle_exceptions,
                )

                for item in generator:
                    if isinstance(item, NodeRunStartedEvent):
                        self.graph_runtime_state.node_run_steps += 1
                        item.route_node_state.index = self.graph_runtime_state.node_run_steps

                    yield item

                self.graph_runtime_state.node_run_state.node_state_mapping[route_node_state.id] = route_node_state

                # append route
                if previous_route_node_state:
                    self.graph_runtime_state.node_run_state.add_route(
                        source_node_state_id=previous_route_node_state.id, target_node_state_id=route_node_state.id
                    )
            except Exception as e:
                route_node_state.status = RouteNodeState.Status.FAILED
                route_node_state.failed_reason = str(e)
                yield NodeRunFailedEvent(
                    error=str(e),
                    id=node_instance.id,
                    node_id=next_node_id,
                    node_type=node_type,
                    node_data=node_instance.node_data,
                    route_node_state=route_node_state,
                    parallel_id=in_parallel_id,
                    parallel_start_node_id=parallel_start_node_id,
                    parent_parallel_id=parent_parallel_id,
                    parent_parallel_start_node_id=parent_parallel_start_node_id,
                )
                raise e

            # It may not be necessary, but it is necessary. :)
            if (
                self.graph.node_id_config_mapping[next_node_id].get("data", {}).get("type", "").lower()
                == NodeType.END.value
            ):
                break

            previous_route_node_state = route_node_state

            # get next node ids
            edge_mappings = self.graph.edge_mapping.get(next_node_id)
            if not edge_mappings:
                break

            if len(edge_mappings) == 1:
                edge = edge_mappings[0]
                if (
                    previous_route_node_state.status == RouteNodeState.Status.EXCEPTION
                    and node_instance.node_data.error_strategy == ErrorStrategy.FAIL_BRANCH
                    and edge.run_condition is None
                ):
                    break
                if edge.run_condition:
                    result = ConditionManager.get_condition_handler(
                        init_params=self.init_params,
                        graph=self.graph,
                        run_condition=edge.run_condition,
                    ).check(
                        graph_runtime_state=self.graph_runtime_state,
                        previous_route_node_state=previous_route_node_state,
                    )

                    if not result:
                        break

                next_node_id = edge.target_node_id
            else:
                final_node_id = None

                if any(edge.run_condition for edge in edge_mappings):
                    # if nodes has run conditions, get node id which branch to take based on the run condition results
                    condition_edge_mappings: dict[str, list[GraphEdge]] = {}
                    for edge in edge_mappings:
                        if edge.run_condition:
                            run_condition_hash = edge.run_condition.hash
                            if run_condition_hash not in condition_edge_mappings:
                                condition_edge_mappings[run_condition_hash] = []

                            condition_edge_mappings[run_condition_hash].append(edge)

                    for _, sub_edge_mappings in condition_edge_mappings.items():
                        if len(sub_edge_mappings) == 0:
                            continue

                        edge = cast(GraphEdge, sub_edge_mappings[0])
                        if edge.run_condition is None:
                            logger.warning(f"Edge {edge.target_node_id} run condition is None")
                            continue

                        result = ConditionManager.get_condition_handler(
                            init_params=self.init_params,
                            graph=self.graph,
                            run_condition=edge.run_condition,
                        ).check(
                            graph_runtime_state=self.graph_runtime_state,
                            previous_route_node_state=previous_route_node_state,
                        )

                        if not result:
                            continue

                        if len(sub_edge_mappings) == 1:
                            final_node_id = edge.target_node_id
                        else:
                            parallel_generator = self._run_parallel_branches(
                                edge_mappings=sub_edge_mappings,
                                in_parallel_id=in_parallel_id,
                                parallel_start_node_id=parallel_start_node_id,
                                handle_exceptions=handle_exceptions,
                            )

                            for parallel_result in parallel_generator:
                                if isinstance(parallel_result, str):
                                    final_node_id = parallel_result
                                else:
                                    yield parallel_result

                        break

                    if not final_node_id:
                        break

                    next_node_id = final_node_id
                elif (
                    node_instance.node_data.error_strategy == ErrorStrategy.FAIL_BRANCH
                    and node_instance.should_continue_on_error
                    and previous_route_node_state.status == RouteNodeState.Status.EXCEPTION
                ):
                    break
                else:
                    parallel_generator = self._run_parallel_branches(
                        edge_mappings=edge_mappings,
                        in_parallel_id=in_parallel_id,
                        parallel_start_node_id=parallel_start_node_id,
                        handle_exceptions=handle_exceptions,
                    )

                    for generated_item in parallel_generator:
                        if isinstance(generated_item, str):
                            final_node_id = generated_item
                        else:
                            yield generated_item

                    if not final_node_id:
                        break

                    next_node_id = final_node_id

            if in_parallel_id and self.graph.node_parallel_mapping.get(next_node_id, "") != in_parallel_id:
                break

    def _run_parallel_branches(
        self,
        edge_mappings: list[GraphEdge],
        in_parallel_id: Optional[str] = None,
        parallel_start_node_id: Optional[str] = None,
        handle_exceptions: list[str] = [],
    ) -> Generator[GraphEngineEvent | str, None, None]:
        # if nodes has no run conditions, parallel run all nodes
        parallel_id = self.graph.node_parallel_mapping.get(edge_mappings[0].target_node_id)
        if not parallel_id:
            node_id = edge_mappings[0].target_node_id
            node_config = self.graph.node_id_config_mapping.get(node_id)
            if not node_config:
                raise GraphRunFailedError(
                    f"Node {node_id} related parallel not found or incorrectly connected to multiple parallel branches."
                )

            node_title = node_config.get("data", {}).get("title")
            raise GraphRunFailedError(
                f"Node {node_title} related parallel not found or incorrectly connected to multiple parallel branches."
            )

        parallel = self.graph.parallel_mapping.get(parallel_id)
        if not parallel:
            raise GraphRunFailedError(f"Parallel {parallel_id} not found.")

        # run parallel nodes, run in new thread and use queue to get results
        q: queue.Queue = queue.Queue()

        # Create a list to store the threads
        futures = []

        # new thread
        for edge in edge_mappings:
            if (
                edge.target_node_id not in self.graph.node_parallel_mapping
                or self.graph.node_parallel_mapping.get(edge.target_node_id, "") != parallel_id
            ):
                continue

            future = self.thread_pool.submit(
                self._run_parallel_node,
                **{
                    "flask_app": current_app._get_current_object(),  # type: ignore[attr-defined]
                    "q": q,
                    "parallel_id": parallel_id,
                    "parallel_start_node_id": edge.target_node_id,
                    "parent_parallel_id": in_parallel_id,
                    "parent_parallel_start_node_id": parallel_start_node_id,
                    "handle_exceptions": handle_exceptions,
                },
            )

            future.add_done_callback(self.thread_pool.task_done_callback)

            futures.append(future)

        succeeded_count = 0
        while True:
            try:
                event = q.get(timeout=1)
                if event is None:
                    break

                yield event
                if event.parallel_id == parallel_id:
                    if isinstance(event, ParallelBranchRunSucceededEvent):
                        succeeded_count += 1
                        if succeeded_count == len(futures):
                            q.put(None)

                        continue
                    elif isinstance(event, ParallelBranchRunFailedEvent):
                        raise GraphRunFailedError(event.error)
            except queue.Empty:
                continue

        # wait all threads
        wait(futures)

        # get final node id
        final_node_id = parallel.end_to_node_id
        if final_node_id:
            yield final_node_id

    def _run_parallel_node(
        self,
        flask_app: Flask,
        q: queue.Queue,
        parallel_id: str,
        parallel_start_node_id: str,
        parent_parallel_id: Optional[str] = None,
        parent_parallel_start_node_id: Optional[str] = None,
        handle_exceptions: list[str] = [],
    ) -> None:
        """
        Run parallel nodes
        """
        with flask_app.app_context():
            try:
                q.put(
                    ParallelBranchRunStartedEvent(
                        parallel_id=parallel_id,
                        parallel_start_node_id=parallel_start_node_id,
                        parent_parallel_id=parent_parallel_id,
                        parent_parallel_start_node_id=parent_parallel_start_node_id,
                    )
                )

                # run node
                generator = self._run(
                    start_node_id=parallel_start_node_id,
                    in_parallel_id=parallel_id,
                    parent_parallel_id=parent_parallel_id,
                    parent_parallel_start_node_id=parent_parallel_start_node_id,
                    handle_exceptions=handle_exceptions,
                )

                for item in generator:
                    q.put(item)

                # trigger graph run success event
                q.put(
                    ParallelBranchRunSucceededEvent(
                        parallel_id=parallel_id,
                        parallel_start_node_id=parallel_start_node_id,
                        parent_parallel_id=parent_parallel_id,
                        parent_parallel_start_node_id=parent_parallel_start_node_id,
                    )
                )
            except GraphRunFailedError as e:
                q.put(
                    ParallelBranchRunFailedEvent(
                        parallel_id=parallel_id,
                        parallel_start_node_id=parallel_start_node_id,
                        parent_parallel_id=parent_parallel_id,
                        parent_parallel_start_node_id=parent_parallel_start_node_id,
                        error=e.error,
                    )
                )
            except Exception as e:
                logger.exception("Unknown Error when generating in parallel")
                q.put(
                    ParallelBranchRunFailedEvent(
                        parallel_id=parallel_id,
                        parallel_start_node_id=parallel_start_node_id,
                        parent_parallel_id=parent_parallel_id,
                        parent_parallel_start_node_id=parent_parallel_start_node_id,
                        error=str(e),
                    )
                )
            finally:
                db.session.remove()

    def _run_node(
        self,
        node_instance: BaseNode[BaseNodeData],
        route_node_state: RouteNodeState,
        parallel_id: Optional[str] = None,
        parallel_start_node_id: Optional[str] = None,
        parent_parallel_id: Optional[str] = None,
        parent_parallel_start_node_id: Optional[str] = None,
        handle_exceptions: list[str] = [],
    ) -> Generator[GraphEngineEvent, None, None]:
        """
        Run node
        """
        # trigger node run start event
        yield NodeRunStartedEvent(
            id=node_instance.id,
            node_id=node_instance.node_id,
            node_type=node_instance.node_type,
            node_data=node_instance.node_data,
            route_node_state=route_node_state,
            predecessor_node_id=node_instance.previous_node_id,
            parallel_id=parallel_id,
            parallel_start_node_id=parallel_start_node_id,
            parent_parallel_id=parent_parallel_id,
            parent_parallel_start_node_id=parent_parallel_start_node_id,
        )

        db.session.close()
        max_retries = node_instance.node_data.retry_config.max_retries
        retry_interval = node_instance.node_data.retry_config.retry_interval_seconds
        retries = 0
        should_continue_retry = True
        while should_continue_retry and retries <= max_retries:
            try:
                # run node
                retry_start_at = datetime.now(UTC).replace(tzinfo=None)
                generator = node_instance.run()
                for item in generator:
                    if isinstance(item, GraphEngineEvent):
                        if isinstance(item, BaseIterationEvent):
                            # add parallel info to iteration event
                            item.parallel_id = parallel_id
                            item.parallel_start_node_id = parallel_start_node_id
                            item.parent_parallel_id = parent_parallel_id
                            item.parent_parallel_start_node_id = parent_parallel_start_node_id

                        yield item
                    else:
                        if isinstance(item, RunCompletedEvent):
                            run_result = item.run_result
                            if run_result.status == WorkflowNodeExecutionStatus.FAILED:
                                if (
                                    retries == max_retries
                                    and node_instance.node_type == NodeType.HTTP_REQUEST
                                    and run_result.outputs
                                    and not node_instance.should_continue_on_error
                                ):
                                    run_result.status = WorkflowNodeExecutionStatus.SUCCEEDED
                                if node_instance.should_retry and retries < max_retries:
                                    retries += 1
                                    route_node_state.node_run_result = run_result
                                    yield NodeRunRetryEvent(
                                        id=node_instance.id,
                                        node_id=node_instance.node_id,
                                        node_type=node_instance.node_type,
                                        node_data=node_instance.node_data,
                                        route_node_state=route_node_state,
                                        predecessor_node_id=node_instance.previous_node_id,
                                        parallel_id=parallel_id,
                                        parallel_start_node_id=parallel_start_node_id,
                                        parent_parallel_id=parent_parallel_id,
                                        parent_parallel_start_node_id=parent_parallel_start_node_id,
                                        error=run_result.error or "Unknown error",
                                        retry_index=retries,
                                        start_at=retry_start_at,
                                    )
                                    time.sleep(retry_interval)
                                    continue
                            route_node_state.set_finished(run_result=run_result)

                            if run_result.status == WorkflowNodeExecutionStatus.FAILED:
                                if node_instance.should_continue_on_error:
                                    # if run failed, handle error
                                    run_result = self._handle_continue_on_error(
                                        node_instance,
                                        item.run_result,
                                        self.graph_runtime_state.variable_pool,
                                        handle_exceptions=handle_exceptions,
                                    )
                                    route_node_state.node_run_result = run_result
                                    route_node_state.status = RouteNodeState.Status.EXCEPTION
                                    if run_result.outputs:
                                        for variable_key, variable_value in run_result.outputs.items():
                                            # append variables to variable pool recursively
                                            self._append_variables_recursively(
                                                node_id=node_instance.node_id,
                                                variable_key_list=[variable_key],
                                                variable_value=variable_value,
                                            )
                                    yield NodeRunExceptionEvent(
                                        error=run_result.error or "System Error",
                                        id=node_instance.id,
                                        node_id=node_instance.node_id,
                                        node_type=node_instance.node_type,
                                        node_data=node_instance.node_data,
                                        route_node_state=route_node_state,
                                        parallel_id=parallel_id,
                                        parallel_start_node_id=parallel_start_node_id,
                                        parent_parallel_id=parent_parallel_id,
                                        parent_parallel_start_node_id=parent_parallel_start_node_id,
                                    )
                                    should_continue_retry = False
                                else:
                                    yield NodeRunFailedEvent(
                                        error=route_node_state.failed_reason or "Unknown error.",
                                        id=node_instance.id,
                                        node_id=node_instance.node_id,
                                        node_type=node_instance.node_type,
                                        node_data=node_instance.node_data,
                                        route_node_state=route_node_state,
                                        parallel_id=parallel_id,
                                        parallel_start_node_id=parallel_start_node_id,
                                        parent_parallel_id=parent_parallel_id,
                                        parent_parallel_start_node_id=parent_parallel_start_node_id,
                                    )
                                should_continue_retry = False
                            elif run_result.status == WorkflowNodeExecutionStatus.SUCCEEDED:
                                if node_instance.should_continue_on_error and self.graph.edge_mapping.get(
                                    node_instance.node_id
                                ):
                                    run_result.edge_source_handle = FailBranchSourceHandle.SUCCESS
                                if run_result.metadata and run_result.metadata.get(NodeRunMetadataKey.TOTAL_TOKENS):
                                    # plus state total_tokens
                                    self.graph_runtime_state.total_tokens += int(
                                        run_result.metadata.get(NodeRunMetadataKey.TOTAL_TOKENS)  # type: ignore[arg-type]
                                    )

                                if run_result.llm_usage:
                                    # use the latest usage
                                    self.graph_runtime_state.llm_usage += run_result.llm_usage

                                # append node output variables to variable pool
                                if run_result.outputs:
                                    for variable_key, variable_value in run_result.outputs.items():
                                        # append variables to variable pool recursively
                                        self._append_variables_recursively(
                                            node_id=node_instance.node_id,
                                            variable_key_list=[variable_key],
                                            variable_value=variable_value,
                                        )

                                # When setting metadata, convert to dict first
                                if not run_result.metadata:
                                    run_result.metadata = {}

                                if parallel_id and parallel_start_node_id:
                                    metadata_dict = dict(run_result.metadata)
                                    metadata_dict[NodeRunMetadataKey.PARALLEL_ID] = parallel_id
                                    metadata_dict[NodeRunMetadataKey.PARALLEL_START_NODE_ID] = parallel_start_node_id
                                    if parent_parallel_id and parent_parallel_start_node_id:
                                        metadata_dict[NodeRunMetadataKey.PARENT_PARALLEL_ID] = parent_parallel_id
                                        metadata_dict[NodeRunMetadataKey.PARENT_PARALLEL_START_NODE_ID] = (
                                            parent_parallel_start_node_id
                                        )
                                    run_result.metadata = metadata_dict

                                yield NodeRunSucceededEvent(
                                    id=node_instance.id,
                                    node_id=node_instance.node_id,
                                    node_type=node_instance.node_type,
                                    node_data=node_instance.node_data,
                                    route_node_state=route_node_state,
                                    parallel_id=parallel_id,
                                    parallel_start_node_id=parallel_start_node_id,
                                    parent_parallel_id=parent_parallel_id,
                                    parent_parallel_start_node_id=parent_parallel_start_node_id,
                                )
                                should_continue_retry = False

                            break
                        elif isinstance(item, RunStreamChunkEvent):
                            yield NodeRunStreamChunkEvent(
                                id=node_instance.id,
                                node_id=node_instance.node_id,
                                node_type=node_instance.node_type,
                                node_data=node_instance.node_data,
                                chunk_content=item.chunk_content,
                                from_variable_selector=item.from_variable_selector,
                                route_node_state=route_node_state,
                                parallel_id=parallel_id,
                                parallel_start_node_id=parallel_start_node_id,
                                parent_parallel_id=parent_parallel_id,
                                parent_parallel_start_node_id=parent_parallel_start_node_id,
                            )
                        elif isinstance(item, RunRetrieverResourceEvent):
                            yield NodeRunRetrieverResourceEvent(
                                id=node_instance.id,
                                node_id=node_instance.node_id,
                                node_type=node_instance.node_type,
                                node_data=node_instance.node_data,
                                retriever_resources=item.retriever_resources,
                                context=item.context,
                                route_node_state=route_node_state,
                                parallel_id=parallel_id,
                                parallel_start_node_id=parallel_start_node_id,
                                parent_parallel_id=parent_parallel_id,
                                parent_parallel_start_node_id=parent_parallel_start_node_id,
                            )
            except GenerateTaskStoppedError:
                # trigger node run failed event
                route_node_state.status = RouteNodeState.Status.FAILED
                route_node_state.failed_reason = "Workflow stopped."
                yield NodeRunFailedEvent(
                    error="Workflow stopped.",
                    id=node_instance.id,
                    node_id=node_instance.node_id,
                    node_type=node_instance.node_type,
                    node_data=node_instance.node_data,
                    route_node_state=route_node_state,
                    parallel_id=parallel_id,
                    parallel_start_node_id=parallel_start_node_id,
                    parent_parallel_id=parent_parallel_id,
                    parent_parallel_start_node_id=parent_parallel_start_node_id,
                )
                return
            except Exception as e:
                logger.exception(f"Node {node_instance.node_data.title} run failed")
                raise e
            finally:
                db.session.close()

    def _append_variables_recursively(self, node_id: str, variable_key_list: list[str], variable_value: VariableValue):
        """
        Append variables recursively
        :param node_id: node id
        :param variable_key_list: variable key list
        :param variable_value: variable value
        :return:
        """
        self.graph_runtime_state.variable_pool.add([node_id] + variable_key_list, variable_value)

        # if variable_value is a dict, then recursively append variables
        if isinstance(variable_value, dict):
            for key, value in variable_value.items():
                # construct new key list
                new_key_list = variable_key_list + [key]
                self._append_variables_recursively(
                    node_id=node_id, variable_key_list=new_key_list, variable_value=value
                )

    def _is_timed_out(self, start_at: float, max_execution_time: int) -> bool:
        """
        Check timeout
        :param start_at: start time
        :param max_execution_time: max execution time
        :return:
        """
        return time.perf_counter() - start_at > max_execution_time

    def create_copy(self):
        """
        create a graph engine copy
        :return: with a new variable pool instance of graph engine
        """
        new_instance = copy(self)
        new_instance.graph_runtime_state = copy(self.graph_runtime_state)
        new_instance.graph_runtime_state.variable_pool = deepcopy(self.graph_runtime_state.variable_pool)
        return new_instance

    def _handle_continue_on_error(
        self,
        node_instance: BaseNode[BaseNodeData],
        error_result: NodeRunResult,
        variable_pool: VariablePool,
        handle_exceptions: list[str] = [],
    ) -> NodeRunResult:
        """
        handle continue on error when self._should_continue_on_error is True


        :param    error_result (NodeRunResult): error run result
        :param    variable_pool (VariablePool): variable pool
        :return:  excption run result
        """
        # add error message and error type to variable pool
        variable_pool.add([node_instance.node_id, "error_message"], error_result.error)
        variable_pool.add([node_instance.node_id, "error_type"], error_result.error_type)
        # add error message to handle_exceptions
        handle_exceptions.append(error_result.error or "")
        node_error_args: dict[str, Any] = {
            "status": WorkflowNodeExecutionStatus.EXCEPTION,
            "error": error_result.error,
            "inputs": error_result.inputs,
            "metadata": {
                NodeRunMetadataKey.ERROR_STRATEGY: node_instance.node_data.error_strategy,
            },
        }

        if node_instance.node_data.error_strategy is ErrorStrategy.DEFAULT_VALUE:
            return NodeRunResult(
                **node_error_args,
                outputs={
                    **node_instance.node_data.default_value_dict,
                    "error_message": error_result.error,
                    "error_type": error_result.error_type,
                },
            )
        elif node_instance.node_data.error_strategy is ErrorStrategy.FAIL_BRANCH:
            if self.graph.edge_mapping.get(node_instance.node_id):
                node_error_args["edge_source_handle"] = FailBranchSourceHandle.FAILED
            return NodeRunResult(
                **node_error_args,
                outputs={
                    "error_message": error_result.error,
                    "error_type": error_result.error_type,
                },
            )
        return error_result


class GraphRunFailedError(Exception):
    def __init__(self, error: str):
        self.error = error
