#!/usr/bin/env python3
"""
工作流验证脚本

验证工作流的结构是否正确，包括：
- 节点类型是否有效
- 连接是否合法
- 必需字段是否存在
- 是否存在循环依赖
"""

import json
import sys
from typing import Dict, List, Set

# 有效的节点类型
VALID_NODE_TYPES = {
    "PROMPT_INPUT",
    "PROMPT_GENERATOR",
    "IMAGE_GENERATOR",
    "VIDEO_GENERATOR",
    "VIDEO_ANALYZER",
    "AUDIO_GENERATOR",
    "INFO_DISPLAY",
    "IMAGE_EDITOR",
}

# 有效的连接规则
VALID_CONNECTIONS = {
    "PROMPT_INPUT": {"PROMPT_GENERATOR", "IMAGE_GENERATOR", "VIDEO_GENERATOR"},
    "PROMPT_GENERATOR": {"IMAGE_GENERATOR", "VIDEO_GENERATOR"},
    "IMAGE_GENERATOR": {"VIDEO_GENERATOR", "IMAGE_EDITOR", "INFO_DISPLAY"},
    "VIDEO_GENERATOR": {"VIDEO_ANALYZER", "INFO_DISPLAY"},
    "VIDEO_ANALYZER": {"INFO_DISPLAY"},
    "AUDIO_GENERATOR": {"INFO_DISPLAY"},
    "IMAGE_EDITOR": {"INFO_DISPLAY"},
}

# 必需字段
REQUIRED_FIELDS = {
    "PROMPT_INPUT": ["prompt", "label"],
    "PROMPT_GENERATOR": ["style", "label"],
    "IMAGE_GENERATOR": ["model", "label"],
    "VIDEO_GENERATOR": ["model", "label"],
}


def validate_workflow(workflow: Dict) -> List[str]:
    """验证工作流
    
    Args:
        workflow: 工作流字典
    
    Returns:
        错误列表（空列表表示验证通过）
    """
    errors = []
    
    # 检查必需字段
    if "nodes" not in workflow:
        errors.append("Missing 'nodes' field")
        return errors
    
    if "edges" not in workflow:
        errors.append("Missing 'edges' field")
        return errors
    
    nodes = workflow["nodes"]
    edges = workflow["edges"]
    
    # 验证节点
    node_ids = set()
    node_types = {}
    
    for i, node in enumerate(nodes):
        # 检查节点 ID
        if "id" not in node:
            errors.append(f"Node {i}: Missing 'id' field")
            continue
        
        node_id = node["id"]
        
        # 检查重复 ID
        if node_id in node_ids:
            errors.append(f"Node {node_id}: Duplicate node ID")
        node_ids.add(node_id)
        
        # 检查节点类型
        if "type" not in node:
            errors.append(f"Node {node_id}: Missing 'type' field")
            continue
        
        node_type = node["type"]
        
        # 检查类型是否有效
        if node_type not in VALID_NODE_TYPES:
            errors.append(f"Node {node_id}: Invalid node type '{node_type}'")
            continue
        
        node_types[node_id] = node_type
        
        # 检查必需字段
        if node_type in REQUIRED_FIELDS:
            if "data" not in node:
                errors.append(f"Node {node_id}: Missing 'data' field")
                continue
            
            data = node["data"]
            for field in REQUIRED_FIELDS[node_type]:
                if field not in data:
                    errors.append(f"Node {node_id}: Missing required field '{field}'")
    
    # 验证连接
    for i, edge in enumerate(edges):
        # 检查必需字段
        if "source" not in edge:
            errors.append(f"Edge {i}: Missing 'source' field")
            continue
        
        if "target" not in edge:
            errors.append(f"Edge {i}: Missing 'target' field")
            continue
        
        source = edge["source"]
        target = edge["target"]
        
        # 检查节点是否存在
        if source not in node_ids:
            errors.append(f"Edge {i}: Source node '{source}' does not exist")
            continue
        
        if target not in node_ids:
            errors.append(f"Edge {i}: Target node '{target}' does not exist")
            continue
        
        # 检查连接是否合法
        source_type = node_types[source]
        target_type = node_types[target]
        
        # INFO_DISPLAY 可以接受任何输入
        if target_type == "INFO_DISPLAY":
            continue
        
        # 检查连接规则
        if source_type in VALID_CONNECTIONS:
            if target_type not in VALID_CONNECTIONS[source_type]:
                errors.append(
                    f"Edge {i}: Invalid connection from {source_type} to {target_type}"
                )
    
    # 检查循环依赖
    if not errors:  # 只在没有其他错误时检查
        cycles = detect_cycles(edges, node_ids)
        if cycles:
            errors.append(f"Circular dependency detected: {' -> '.join(cycles)}")
    
    return errors


def detect_cycles(edges: List[Dict], node_ids: Set[str]) -> List[str]:
    """检测循环依赖
    
    Args:
        edges: 边列表
        node_ids: 节点 ID 集合
    
    Returns:
        循环路径（空列表表示无循环）
    """
    # 构建邻接表
    graph = {node_id: [] for node_id in node_ids}
    for edge in edges:
        if "source" in edge and "target" in edge:
            graph[edge["source"]].append(edge["target"])
    
    # DFS 检测循环
    visited = set()
    rec_stack = set()
    path = []
    
    def dfs(node):
        visited.add(node)
        rec_stack.add(node)
        path.append(node)
        
        for neighbor in graph[node]:
            if neighbor not in visited:
                if dfs(neighbor):
                    return True
            elif neighbor in rec_stack:
                # 找到循环
                cycle_start = path.index(neighbor)
                return path[cycle_start:]
        
        path.pop()
        rec_stack.remove(node)
        return False
    
    for node in node_ids:
        if node not in visited:
            result = dfs(node)
            if result:
                return result
    
    return []


def main():
    """主函数"""
    if len(sys.argv) < 2:
        print("Usage: validate_workflow.py <workflow.json>")
        sys.exit(1)
    
    workflow_file = sys.argv[1]
    
    try:
        with open(workflow_file, 'r', encoding='utf-8') as f:
            workflow = json.load(f)
    except FileNotFoundError:
        print(f"Error: File '{workflow_file}' not found")
        sys.exit(1)
    except json.JSONDecodeError as e:
        print(f"Error: Invalid JSON - {e}")
        sys.exit(1)
    
    # 验证工作流
    errors = validate_workflow(workflow)
    
    if errors:
        print("❌ Validation failed:")
        for error in errors:
            print(f"  - {error}")
        sys.exit(1)
    else:
        print("✅ Workflow validation passed")
        sys.exit(0)


if __name__ == "__main__":
    main()
