from .LiblibUtil import LiblibUtil
import time


def get_generation_status(generate_uuid):
    """查询生图任务状态和结果"""
    liblib_util = LiblibUtil()
    uri = "/api/generate/webui/status"
    payload = {"generateUuid": generate_uuid}
    return liblib_util.post_request(uri, payload)


def wait_for_generation_completion(generate_uuid, check_interval=5, max_wait_time=300):
    """等待生图任务完成并获取最终结果"""
    liblib_util = LiblibUtil()
    start_time = time.time()
    
    while time.time() - start_time < max_wait_time:
        status_data = get_generation_status(generate_uuid)
        
        if status_data:
            status = status_data.get('status')
            progress = status_data.get('percentCompleted', 0) * 100
            print(f"生成进度：{progress:.2f}%")
            
            if status == 'completed':
                print("生图任务已完成")
                images = status_data.get('images', [])
                if images:
                    print("生成的最终图片：")
                    for img in images:
                        print(f"- {img.get('imageUrl')}")
                return status_data
            elif status == 'failed':
                print(f"生图任务失败: {status_data.get('errorMessage', '未知错误')}")
                return None
            else:
                print(f"生图任务处理中，当前状态: {status}")
        
        time.sleep(check_interval)
    
    print(f"生图任务超时（{max_wait_time}秒）")
    return None


if __name__ == '__main__':
    print("===== 测试查询生图结果 =====")
    test_generate_uuid = "061d402be5af492d9949d53719912737"
    
    if test_generate_uuid:
        status_data = get_generation_status(test_generate_uuid)
        if status_data:
            print(f"生图任务状态：{status_data}")
        else:
            print("查询生图任务状态失败")
    else:
        print("请先设置有效的生图任务UUID")