"""基础用法示例 - RoboDock SDK."""

import logging
from robodock_sdk import RoboDockClient
from robodock_sdk.actions import RoboDockActions

# 配置日志
logging.basicConfig(level=logging.INFO)

def main():
    """基础SDK使用示例。"""
    
    # 1. 创建客户端
    print("=" * 50)
    print("1. 创建RoboDock客户端")
    print("=" * 50)
    
    client = RoboDockClient(
        s3_endpoint="http://192.168.1.6:9000",
        s3_access_key="robodock",
        s3_secret_key="robodock",
        api_url="http://192.168.1.2:31080"
    )
    print("✓ 客户端创建成功\n")
    
    # 2. 获取数据集信息
    print("=" * 50)
    print("2. 获取数据集信息")
    print("=" * 50)
    
    dataset_uuid = "your-dataset-uuid"
    try:
        dataset_info = client.get_dataset(dataset_uuid)
        print(f"数据集UUID: {dataset_uuid}")
        print(f"数据集信息: {dataset_info}\n")
    except Exception as e:
        print(f"获取数据集信息失败: {e}\n")
    
    # 3. 获取单条数据信息
    print("=" * 50)
    print("3. 获取单条数据信息")
    print("=" * 50)
    
    data_uuid = "your-data-uuid"
    try:
        data_info = client.get_data(data_uuid)
        print(f"数据UUID: {data_uuid}")
        print(f"数据信息: {data_info}\n")
    except Exception as e:
        print(f"获取数据信息失败: {e}\n")
    
    # 4. 下载数据集
    print("=" * 50)
    print("4. 下载数据集")
    print("=" * 50)
    
    try:
        local_path = client.download_dataset(
            data_name="example_data.zip",
            uuid=data_uuid,
            dirname="singledata"
        )
        print(f"数据已下载到: {local_path}\n")
    except Exception as e:
        print(f"下载数据失败: {e}\n")
    
    # 5. 获取运行任务信息
    print("=" * 50)
    print("5. 获取运行任务信息")
    print("=" * 50)
    
    run_name = "project-run001"
    try:
        run_info = client.get_run(run_name, unique_type="unique_name")
        print(f"运行名称: {run_name}")
        print(f"运行信息: {run_info}\n")
    except Exception as e:
        print(f"获取运行信息失败: {e}\n")
    
    # 6. 更新任务状态
    print("=" * 50)
    print("6. 更新任务状态")
    print("=" * 50)
    
    import time
    try:
        # 更新为处理中状态
        client.update_run_status(
            unique_name=run_name,
            status="processing",
            start_time=int(time.time() * 1000)
        )
        print(f"任务状态已更新为: processing\n")
        
        # 更新为成功状态
        client.update_run_status(
            unique_name=run_name,
            status="success",
            finish_time=int(time.time() * 1000)
        )
        print(f"任务状态已更新为: success\n")
    except Exception as e:
        print(f"更新任务状态失败: {e}\n")
    
    # 7. 上传模型
    print("=" * 50)
    print("7. 上传模型")
    print("=" * 50)
    
    try:
        s3_key = client.upload_model(
            outdir_path="/path/to/your/model",
            upload_path="models/project/run001"
        )
        print(f"模型已上传至: {s3_key}\n")
    except Exception as e:
        print(f"上传模型失败: {e}\n")

    # 8. 异步上传Checkpoint
    print("=" * 50)
    print("8. 异步上传Checkpoint")
    print("=" * 50)

    checkpoint_dir = "/path/to/checkpoint_dir"
    checkpoint_upload_path = "checkpoints/project/run001"

    try:
        future = client.upload_checkpoint_async(
            outdir_path=checkpoint_dir,
            upload_path=checkpoint_upload_path,
            run_unique_name=run_name,
            unique_type="unique_name",
            run_field="check_point_path",
        )
        if future is None:
            print("上传队列已满，本次上传被跳过。\n")
        else:
            print("已提交异步上传任务，主流程继续执行。\n")
            RoboDockActions.wait_for_uploads()
            print("所有异步上传任务已完成。\n")
    except Exception as e:
        print(f"异步上传失败: {e}\n")


if __name__ == "__main__":
    main()
