"""
TraceStudio 节点图运行时数据结构设计

本文件定义了节点图执行所需的核心数据结构，确保前后端一致。
"""
from enum import Enum
from typing import Any, Dict, List, Optional
from dataclasses import dataclass, field
from app.core.node_base import NodeType

class DimensionMode(Enum):
    NONE = "none"
    EXPAND = "expand"
    COLLAPSE = "collapse"
    BROADCAST = "broadcast"

@dataclass
class EdgeMetadata:
    source_node: str
    source_port: str
    target_node: str
    target_port: str
    dimension_mode: DimensionMode = DimensionMode.NONE

@dataclass
class NodeMetadata:
    node_id: str
    # 严格约定：type 必须为 NodeType 的值（"input"/"normal"/"output"/"function"/...）
    type: str
    # 实现类名（必须提供，用于从 NodeRegistry 加载实现）
    class_name: str
    params: Dict[str, Any] = field(default_factory=dict)
    sub_workflow_nodes: Optional[List[Dict]] = None
    sub_workflow_edges: Optional[List[Dict]] = None

@dataclass
class GraphRuntime:
    nodes: List[NodeMetadata]
    edges: List[EdgeMetadata]
    # 可扩展：全局上下文、运行参数等

    def to_executor_format(self) -> Dict[str, Any]:
        """
        将运行时 Graph 格式转换为 WorkflowExecutor 接受的 nodes/edges 列表。

        - 将 EdgeMetadata 字段命名转换为 executor 期望的 snake_case（source, source_port, target, target_port）
        - 将 NodeMetadata 中的 sub_workflow_* 字段展开为 executor 可识别的键名（sub_workflow_nodes/sub_workflow_edges）
        - 保持 params 原样传递
        """
        nodes_out = []
        for n in self.nodes:
            # 验证 type 为 NodeType 的合法值
            try:
                NodeType(n.type)
            except Exception:
                raise ValueError(f"NodeMetadata.type 必须为 NodeType 的值之一: {[t.value for t in NodeType]} (got {n.type})")

            nd = {
                "id": n.node_id,
                "type": n.type,
                "class_name": n.class_name,
                "params": n.params or {}
            }
            # 支持两种约定：直接 sub_workflow_* 或嵌套 sub_workflow
            if n.sub_workflow_nodes:
                nd["sub_workflow_nodes"] = n.sub_workflow_nodes
            if n.sub_workflow_edges:
                nd["sub_workflow_edges"] = n.sub_workflow_edges
            nodes_out.append(nd)

        edges_out = []
        for e in self.edges:
            edges_out.append({
                "source": e.source_node,
                "source_port": e.source_port,
                "target": e.target_node,
                "target_port": e.target_port,
                # executor currently ignores dimension_mode on input dicts,
                # but we include it for completeness (as string)
                "dimension_mode": e.dimension_mode.value if e.dimension_mode else DimensionMode.NONE.value
            })

        return {"nodes": nodes_out, "edges": edges_out}