from threading import Lock
from typing import Any, Optional


class VariablePool:
    """
    变量池 - 工作流全局上下文容器

    设计要点：
    1. 命名空间隔离：每个变量由 (namespace, key) 二元组唯一标识
       - 系统内置命名空间：'sys'（系统变量）、'env'（环境变量）、'session'（会话变量）
       - 节点输出命名空间：以 node_id 为 namespace，节点执行后写入
    2. 线程安全：写操作（set）通过锁保护；读操作在写锁释放后进行
    3. 不可变原则：已写入的变量值不允许覆盖（同一节点同一变量名只能写一次）
                   防止重试或并发导致的数据错乱
    """

    def __init__(
        self,
        user_inputs: dict[str, Any],
        system_variables: dict[str, Any],
        environment_variables: dict[str, Any],
        session_variables: dict[str, Any],
    ) -> None:
        self._lock = Lock()

        # 内部存储：双层字典，外层 Key 为命名空间，内层 Key 为变量名
        # 使用内置命名空间预填充
        self._pool: dict[str, dict[str, Any]] = {
            "sys":     system_variables,       # 系统级变量（workflow_id、user_id 等）
            "env":     environment_variables,  # 环境变量（API Key、数据库地址等敏感配置）
            "session": session_variables,      # 会话变量（跨轮次对话保留的上下文）
            "user":    user_inputs,            # 用户输入（本次触发时用户提供的参数）
        }

    def set(self, node_id: str, key: str, value: Any) -> None:
        """
        写入节点输出变量

        - 以 node_id 作为命名空间，确保不同节点的输出互不干扰
        - 幂等保护：已存在的变量不允许覆盖，防止重试场景下的数据污染
        """
        with self._lock:
            if node_id not in self._pool:
                self._pool[node_id] = {}

            if key in self._pool[node_id]:
                # 已存在则忽略（幂等），避免重试时覆盖正确结果
                return

            self._pool[node_id][key] = value

    def get(self, variable_selector: list[str]) -> Optional[Any]:
        """
        读取变量

        variable_selector: 变量路径，格式为 [namespace, key]
        例如：
            ["user", "query"]          → 用户输入的 query 字段
            ["sys", "workflow_run_id"] → 系统变量
            ["llm_node_1", "output"]   → llm_node_1 节点的 output 输出

        读操作不需要加锁（Python GIL 保证基本读操作原子性），
        但为了保障强一致性语义，此处仍加锁
        """
        if len(variable_selector) < 2:
            raise ValueError(f"变量选择器格式错误: {variable_selector}，至少需要 [namespace, key]")

        namespace = variable_selector[0]
        key = variable_selector[1]

        with self._lock:
            namespace_dict = self._pool.get(namespace)
            if namespace_dict is None:
                return None
            return namespace_dict.get(key)

    def get_or_raise(self, variable_selector: list[str]) -> Any:
        """读取变量，不存在则抛出异常（用于必填输入的校验）"""
        value = self.get(variable_selector)
        if value is None:
            raise KeyError(
                f"变量 {variable_selector} 在变量池中不存在。"
                f"请检查上游节点是否已成功执行，或变量名是否正确。"
            )
        return value

    def has(self, variable_selector: list[str]) -> bool:
        """检查变量是否存在"""
        return self.get(variable_selector) is not None

import re
from typing import Any


class VariableResolver:
    """
    变量解析器 - 负责将节点配置中的变量引用替换为实际值

    支持两种引用方式：
    1. 模板语法（用于字符串拼接）：{{ node_id.key }}
       例如："请基于以下内容回答：{{ retrieval_node.documents }}"

    2. 变量选择器（用于直接传值）：["node_id", "key"]
       例如：inputs["query"] = ["start_node", "user_query"]
              → inputs["query"] = "用户实际输入的问题"
    """

    # 匹配 {{ namespace.key }} 格式（支持多级路径，如 {{ a.b.c }}）
    TEMPLATE_PATTERN = re.compile(r'\{\{\s*([\w]+\.[\w.]+)\s*\}\}')

    def __init__(self, variable_pool: VariablePool) -> None:
        self._pool = variable_pool

    def resolve_template(self, template: str) -> str:
        """
        解析模板字符串，将所有 {{ namespace.key }} 替换为实际值

        示例：
            template = "用户问：{{ user.query }}，参考：{{ retrieval_node.context }}"
            → "用户问：什么是DAG？，参考：DAG是有向无环图..."
        """
        def replace_match(match: re.Match) -> str:
            # "retrieval_node.context" → ["retrieval_node", "context"]
            parts = match.group(1).split('.')
            value = self._pool.get(parts)
            if value is None:
                raise KeyError(
                    f"模板变量 '{{{{ {match.group(1)} }}}}' 无法解析：变量不存在于变量池中。"
                    f"请确认上游节点已成功执行且输出了该变量。"
                )
            # 列表和字典转为字符串表示，便于拼入模板
            if isinstance(value, (list, dict)):
                return str(value)
            return str(value)

        return self.TEMPLATE_PATTERN.sub(replace_match, template)

    def resolve_selector(self, selector: list[str]) -> Any:
        """
        解析变量选择器，返回实际值

        selector: ["namespace", "key"]，如 ["llm_node", "reply"]
        """
        return self._pool.get_or_raise(selector)

    def resolve_node_inputs(self, input_config: dict[str, Any]) -> dict[str, Any]:
        """
        批量解析节点的所有输入配置

        - 字符串类型的值：尝试作为模板解析
        - 列表类型的值（变量选择器格式）：直接解析为实际值
        - 其他类型：原样传递

        示例：
            input_config = {
                "prompt":    "请基于{{ retrieval_node.documents }}回答{{ user.query }}",
                "max_tokens": 512,
                "context":   ["retrieval_node", "documents"],
            }
            →  resolved = {
                "prompt":    "请基于[检索到的文档...]回答什么是DAG？",
                "max_tokens": 512,
                "context":   ["DAG是有向无环图...", "Kahn算法用于拓扑排序..."],
            }
        """
        resolved: dict[str, Any] = {}

        for key, value in input_config.items():
            if isinstance(value, str):
                # 字符串：作为模板解析（支持 {{ }} 语法）
                resolved[key] = self.resolve_template(value)

            elif (isinstance(value, list)
                  and len(value) == 2
                  and all(isinstance(v, str) for v in value)):
                # 二元字符串列表：视为变量选择器 ["namespace", "key"]
                resolved[key] = self.resolve_selector(value)

            elif isinstance(value, dict):
                # 嵌套字典：递归解析
                resolved[key] = self.resolve_node_inputs(value)

            else:
                # 数字、布尔值等：原样传递
                resolved[key] = value

        return resolved

from enum import Enum
from typing import Set


class NodeRunStatus(str, Enum):
    """
    节点运行状态枚举 - 状态机的"状态"定义

    继承 str 使其可以直接序列化为 JSON，便于持久化存储
    """
    PENDING   = "pending"    # 等待中：前置依赖尚未全部完成
    RUNNING   = "running"    # 运行中：正在线程池中执行
    SUCCEEDED = "succeeded"  # 已成功：执行完毕，输出已写入变量池
    FAILED    = "failed"     # 已失败：执行过程中抛出异常（终态）
    SKIPPED   = "skipped"    # 已跳过：条件分支不满足，未执行（终态）


# 状态机的"转换函数" - 定义每个状态允许跃迁到的目标状态集合
# 这是整个状态机的核心约束，不在此表中的跃迁将被直接拒绝
VALID_TRANSITIONS: dict[NodeRunStatus, Set[NodeRunStatus]] = {
    NodeRunStatus.PENDING: {
        NodeRunStatus.RUNNING,   # 正常调度：前置依赖满足，开始执行
        NodeRunStatus.SKIPPED,   # 条件分支：当前节点所在分支条件不满足
    },
    NodeRunStatus.RUNNING: {
        NodeRunStatus.SUCCEEDED, # 正常完成：节点逻辑执行无误
        NodeRunStatus.FAILED,    # 执行异常：节点抛出未捕获异常
    },
    # 终态：不允许任何跃迁
    NodeRunStatus.SUCCEEDED: set(),
    NodeRunStatus.FAILED:    set(),
    NodeRunStatus.SKIPPED:   set(),
}

import time
from dataclasses import dataclass, field
from typing import Any, Optional


@dataclass
class NodeExecutionStats:
    """
    单个节点的运行时统计信息 - 节点状态机的"数据载体"

    记录节点从 PENDING 到终态的完整执行轨迹，包括：
    - 状态（State）：当前处于哪个生命周期阶段
    - 时间戳（Timestamps）：各阶段的精确时间，用于性能分析
    - 数据快照（Snapshots）：输入/输出的完整记录，用于问题追溯
    """
    node_id:   str
    node_type: str

    # 当前状态，初始为 PENDING
    status: NodeRunStatus = NodeRunStatus.PENDING

    # 时间戳
    started_at:  Optional[float] = None  # 开始执行时间（RUNNING 时记录）
    finished_at: Optional[float] = None  # 执行结束时间（终态时记录）

    # 数据快照 - 用于全链路追溯
    inputs:  dict[str, Any] = field(default_factory=dict)  # 节点输入（变量解析后）
    outputs: dict[str, Any] = field(default_factory=dict)  # 节点输出（写入变量池前）

    # 错误信息（仅 FAILED 状态时有值）
    error:       Optional[str] = None
    error_detail: Optional[str] = None  # 完整异常堆栈，用于开发者调试

    @property
    def elapsed_time(self) -> Optional[float]:
        """节点执行耗时（秒），仅在终态时有值"""
        if self.started_at and self.finished_at:
            return self.finished_at - self.started_at
        return None

    @property
    def is_terminal(self) -> bool:
        """是否处于终态（Succeeded / Failed / Skipped）"""
        return self.status in (
            NodeRunStatus.SUCCEEDED,
            NodeRunStatus.FAILED,
            NodeRunStatus.SKIPPED,
        )

import time
from threading import Lock
from typing import Optional


class GraphRuntimeState:
    """
    图运行时状态 - 工作流执行的"总控台"

    职责：
    - 集中管理所有节点的执行状态（状态机入口）
    - 持有变量池的引用（上下文传递的入口）
    - 记录全局执行统计（可观测性）
    - 所有操作线程安全，支持多节点并行执行
    """

    def __init__(self, variable_pool: "VariablePool") -> None:
        # 变量池 - 上下文传递的核心载体
        self.variable_pool = variable_pool

        # 节点状态表 - Key: node_id, Value: NodeExecutionStats
        # 通过 _lock 保护，确保并发安全
        self._node_stats: dict[str, NodeExecutionStats] = {}
        self._lock = Lock()

        # 全局统计
        self.total_steps: int = 0          # 已执行（RUNNING）的节点总数
        self.started_at: float = time.time()

    def initialize_node(self, node_id: str, node_type: str) -> None:
        """
        初始化节点状态为 PENDING

        在 DAG 图构建完成后、执行开始前调用，为所有节点预建状态槽位
        """
        with self._lock:
            self._node_stats[node_id] = NodeExecutionStats(
                node_id=node_id,
                node_type=node_type,
                status=NodeRunStatus.PENDING,
            )

    def transit(
        self,
        node_id: str,
        new_status: NodeRunStatus,
        *,
        inputs: Optional[dict] = None,
        outputs: Optional[dict] = None,
        error: Optional[str] = None,
        error_detail: Optional[str] = None,
    ) -> NodeExecutionStats:
        """
        节点状态转换 - 状态机的核心方法

        严格校验转换合法性：
        - 非法跃迁（如 SUCCEEDED → RUNNING）直接抛出异常，绝不静默放行
        - 合法跃迁自动更新时间戳、数据快照，并维护全局计数器

        并发安全：整个"校验 + 写入"过程在同一个锁范围内原子完成
        """
        with self._lock:
            stats = self._node_stats.get(node_id)
            if stats is None:
                raise ValueError(f"节点 '{node_id}' 未初始化，请先调用 initialize_node()")

            current_status = stats.status
            allowed = VALID_TRANSITIONS.get(current_status, set())

            if new_status not in allowed:
                raise InvalidStateTransitionError(
                    f"节点 '{node_id}': 非法状态跃迁 {current_status} → {new_status}，"
                    f"当前状态允许跃迁到: {allowed}"
                )

            # 更新状态
            stats.status = new_status

            # 根据目标状态，记录相应的时间戳和数据
            if new_status == NodeRunStatus.RUNNING:
                stats.started_at = time.time()
                self.total_steps += 1
                if inputs is not None:
                    stats.inputs = inputs   # 记录解析后的输入快照

            elif new_status in (NodeRunStatus.SUCCEEDED, NodeRunStatus.FAILED, NodeRunStatus.SKIPPED):
                stats.finished_at = time.time()
                if outputs is not None:
                    stats.outputs = outputs
                if error is not None:
                    stats.error = error
                if error_detail is not None:
                    stats.error_detail = error_detail

            return stats

    def get_node_status(self, node_id: str) -> Optional[NodeRunStatus]:
        """线程安全地查询节点当前状态"""
        with self._lock:
            stats = self._node_stats.get(node_id)
            return stats.status if stats else None

    def get_all_stats(self) -> dict[str, NodeExecutionStats]:
        """获取所有节点统计快照（浅拷贝），用于日志或持久化"""
        with self._lock:
            return dict(self._node_stats)


class InvalidStateTransitionError(Exception):
    """非法状态跃迁异常 - 表明引擎内部逻辑出现了严重错误"""
    pass

from abc import ABC
from dataclasses import dataclass, field
from typing import Any, Optional
import time


class WorkflowEvent(ABC):
    """
    工作流事件基类 - 事件驱动架构的基础

    作用：将引擎内部的状态变更解耦地广播给外部观察者：
    - 数据库持久化层：将事件写入数据库，支持断点续跑
    - WebSocket 推送层：将事件实时推送给前端，驱动可视化进度更新
    - 日志系统：将事件写入日志，支持离线分析
    - 监控告警：特定事件（如 NodeFailedEvent）触发告警通知

    设计原则：外部观察者只需消费事件流，无需了解引擎内部逻辑
    """
    pass


@dataclass
class NodeStartedEvent(WorkflowEvent):
    """节点开始执行事件（状态 PENDING → RUNNING 时触发）"""
    node_id:    str
    node_type:  str
    inputs:     dict[str, Any]         # 已解析的输入值快照
    started_at: float = field(default_factory=time.time)


@dataclass
class NodeSucceededEvent(WorkflowEvent):
    """节点执行成功事件（状态 RUNNING → SUCCEEDED 时触发）"""
    node_id:      str
    node_type:    str
    outputs:      dict[str, Any]       # 节点输出值
    elapsed_time: float                # 执行耗时（秒）
    finished_at:  float = field(default_factory=time.time)


@dataclass
class NodeFailedEvent(WorkflowEvent):
    """节点执行失败事件（状态 RUNNING → FAILED 时触发）"""
    node_id:      str
    node_type:    str
    error:        str                  # 错误摘要（面向用户）
    error_detail: Optional[str]        # 异常堆栈（面向开发者）
    elapsed_time: float
    finished_at:  float = field(default_factory=time.time)


@dataclass
class NodeSkippedEvent(WorkflowEvent):
    """节点跳过事件（状态 PENDING → SKIPPED 时触发）"""
    node_id:   str
    node_type: str
    reason:    str                     # 跳过原因（如"条件分支未命中"）


@dataclass
class WorkflowStartedEvent(WorkflowEvent):
    """工作流开始执行事件"""
    workflow_run_id: str
    started_at:      float = field(default_factory=time.time)


@dataclass
class WorkflowCompletedEvent(WorkflowEvent):
    """工作流执行完成事件（所有节点终态）"""
    workflow_run_id: str
    outputs:         dict[str, Any]    # 工作流最终输出（来自 END 节点）
    total_steps:     int               # 实际执行的节点总数
    elapsed_time:    float             # 总耗时
    finished_at:     float = field(default_factory=time.time)


@dataclass
class WorkflowFailedEvent(WorkflowEvent):
    """工作流执行失败事件（任一关键节点失败且无错误处理分支）"""
    workflow_run_id: str
    error:           str
    finished_at:     float = field(default_factory=time.time)