#!/usr/bin/env python3
"""
快速测试脚本 - 验证节点是否可以正常导入和使用

这个脚本可以快速检查：
1. 节点代码是否可以正常导入
2. 节点类是否可以实例化
3. 节点映射是否正确
"""

def quick_test():
    """快速测试节点基本功能"""
    print("🔍 开始快速测试...\n")

    try:
        # 1. 测试导入
        print("1. 测试模块导入...")
        from sora_node import (
            SoraClient,
            SoraTextToVideo,
            SoraImageToVideo,
            SoraVideoToVideo,
            NODE_CLASS_MAPPINGS,
            NODE_DISPLAY_NAME_MAPPINGS
        )
        print("   ✓ 模块导入成功\n")

        # 2. 测试节点类
        print("2. 测试节点类...")
        # Text-to-Video
        node_t2v = SoraTextToVideo()
        input_types = node_t2v.INPUT_TYPES()
        assert "required" in input_types
        assert "prompt" in input_types["required"]
        print("   ✓ SoraTextToVideo 节点正常")
        
        # Image-to-Video
        node_i2v = SoraImageToVideo()
        input_types_i2v = node_i2v.INPUT_TYPES()
        assert "image" in input_types_i2v["required"]
        assert "prompt" in input_types_i2v["required"]
        print("   ✓ SoraImageToVideo 节点正常")
        
        # Video-to-Video
        node_v2v = SoraVideoToVideo()
        input_types_v2v = node_v2v.INPUT_TYPES()
        assert "video" in input_types_v2v["required"]
        assert "prompt" in input_types_v2v["required"]
        print("   ✓ SoraVideoToVideo 节点正常")
        print("   ✓ 所有节点类实例化成功\n")

        # 3. 测试节点映射
        print("3. 测试节点映射...")
        assert "SoraTextToVideo" in NODE_CLASS_MAPPINGS
        assert "SoraImageToVideo" in NODE_CLASS_MAPPINGS
        assert "SoraVideoToVideo" in NODE_CLASS_MAPPINGS
        assert NODE_CLASS_MAPPINGS["SoraTextToVideo"] == SoraTextToVideo
        assert NODE_CLASS_MAPPINGS["SoraImageToVideo"] == SoraImageToVideo
        assert NODE_CLASS_MAPPINGS["SoraVideoToVideo"] == SoraVideoToVideo
        assert "SoraTextToVideo" in NODE_DISPLAY_NAME_MAPPINGS
        assert "SoraImageToVideo" in NODE_DISPLAY_NAME_MAPPINGS
        assert "SoraVideoToVideo" in NODE_DISPLAY_NAME_MAPPINGS
        print("   ✓ 所有节点映射正确\n")

        # 4. 测试客户端初始化（不需要真实 API）
        print("4. 测试客户端类...")
        try:
            client = SoraClient(api_key="test-key-123")
            assert client.api_key == "test-key-123"
            print("   ✓ 客户端类正常\n")
        except Exception as e:
            print(f"   ✗ 客户端测试失败: {e}\n")

        print("=" * 50)
        print("✅ 所有快速测试通过！")
        print("=" * 50)
        print("\n节点已准备好，可以：")
        print("  1. 复制到 ComfyUI/custom_nodes/ 目录")
        print("  2. 重启 ComfyUI")
        print("  3. 在节点菜单中找到以下节点：")
        print("     - Sora Text-To-Video")
        print("     - Sora Image-To-Video")
        print("     - Sora Video-To-Video")
        return True

    except ImportError as e:
        print(f"   ✗ 导入失败: {e}")
        print("\n请确保 sora_node.py 在同一目录下")
        return False
    except Exception as e:
        print(f"   ✗ 测试失败: {e}")
        import traceback
        traceback.print_exc()
        return False


if __name__ == "__main__":
    quick_test()