#!/usr/bin/env python3
"""
测试 ComfyUI 节点加载
模拟 ComfyUI 的加载过程，验证节点是否可以正确导入
"""

import sys
import os

def test_comfyui_load():
    """测试节点是否可以像 ComfyUI 那样加载"""
    print("=" * 60)
    print("测试 ComfyUI 节点加载...")
    print("=" * 60)
    
    # 获取当前目录
    current_dir = os.path.dirname(os.path.abspath(__file__))
    sys.path.insert(0, current_dir)
    
    try:
        # 尝试导入（模拟 ComfyUI 的导入方式）
        print("\n1. 尝试导入节点模块...")
        from __init__ import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
        print("   ✅ 导入成功！")
        
        # 验证节点映射
        print("\n2. 验证节点映射...")
        if not NODE_CLASS_MAPPINGS:
            print("   ❌ NODE_CLASS_MAPPINGS 为空")
            return False
        
        if not NODE_DISPLAY_NAME_MAPPINGS:
            print("   ❌ NODE_DISPLAY_NAME_MAPPINGS 为空")
            return False
        
        print(f"   ✅ 找到 {len(NODE_CLASS_MAPPINGS)} 个节点")
        
        # 列出所有节点
        print("\n3. 已注册的节点:")
        for node_name, node_class in NODE_CLASS_MAPPINGS.items():
            display_name = NODE_DISPLAY_NAME_MAPPINGS.get(node_name, node_name)
            print(f"   - {node_name} ({display_name})")
            print(f"     类: {node_class.__name__}")
        
        # 验证节点类
        print("\n4. 验证节点类...")
        for node_name, node_class in NODE_CLASS_MAPPINGS.items():
            try:
                # 检查必需的属性
                assert hasattr(node_class, 'CATEGORY'), f"{node_name} 缺少 CATEGORY"
                assert hasattr(node_class, 'RETURN_TYPES'), f"{node_name} 缺少 RETURN_TYPES"
                assert hasattr(node_class, 'FUNCTION'), f"{node_name} 缺少 FUNCTION"
                assert hasattr(node_class, 'INPUT_TYPES'), f"{node_name} 缺少 INPUT_TYPES"
                
                # 检查 INPUT_TYPES 方法
                input_types = node_class.INPUT_TYPES()
                assert isinstance(input_types, dict), f"{node_name} INPUT_TYPES 必须返回字典"
                assert "required" in input_types or "optional" in input_types, \
                    f"{node_name} INPUT_TYPES 必须包含 'required' 或 'optional'"
                
                print(f"   ✅ {node_name} 验证通过")
            except Exception as e:
                print(f"   ❌ {node_name} 验证失败: {e}")
                return False
        
        print("\n" + "=" * 60)
        print("✅ 所有测试通过！节点应该可以在 ComfyUI 中正常显示。")
        print("=" * 60)
        print("\n如果 ComfyUI 中仍然不显示节点，请：")
        print("1. 确认文件已复制到 ComfyUI/custom_nodes/sora/ 目录")
        print("2. 完全重启 ComfyUI（不是刷新页面）")
        print("3. 检查 ComfyUI 启动日志中的错误信息")
        return True
        
    except ImportError as e:
        print(f"\n❌ 导入失败: {e}")
        print("\n可能的原因：")
        print("1. 缺少 __init__.py 文件")
        print("2. sora_node.py 文件不存在或路径错误")
        print("3. Python 语法错误")
        import traceback
        traceback.print_exc()
        return False
        
    except Exception as e:
        print(f"\n❌ 测试失败: {e}")
        import traceback
        traceback.print_exc()
        return False


if __name__ == "__main__":
    success = test_comfyui_load()
    sys.exit(0 if success else 1)