import sys
import os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from liblibai_client import LibLibAIClient

def test_generate():
    """测试生图功能"""
    print("开始测试 LibLib ControlNet 生图功能...")
    
    # 使用一个公开的测试图片 URL
    test_image_url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg"
    
    try:
        client = LibLibAIClient()
        print(f"✓ 客户端初始化成功")
        
        # 测试生图
        result = client.generate_image(
            image=test_image_url,
            prompt="beautiful landscape, mountains, blue sky, high quality",
            negative_prompt="lowres, bad anatomy, text, error",
            width=512,
            height=512,
            steps=20,
            cfg_scale=7,
            img_count=1,
            control_weight=1.0,
            preprocessor=1,
            canny_low=100,
            canny_high=200
        )
        
        print(f"✓ 生图任务完成")
        print(f"  任务 ID: {result['task_id']}")
        print(f"  状态: {result['status']}")
        print(f"  生成图片数量: {len(result['images'])}")
        
        if result['images']:
            print(f"  图片 URL:")
            for i, img_url in enumerate(result['images'], 1):
                print(f"    {i}. {img_url}")
        
        assert result['status'] in ['success', 'timeout'], f"任务状态异常: {result['status']}"
        print("\n✓ 测试通过！")
        
    except Exception as e:
        print(f"\n✗ 测试失败: {str(e)}")
        import traceback
        traceback.print_exc()
        sys.exit(1)

if __name__ == "__main__":
    test_generate()