"""完整训练工作流示例 - RoboDock SDK."""

import os
import logging
import time
import shutil
from pathlib import Path
from robodock_sdk import RoboDockClient
from robodock_sdk.utils import get_env_config, is_resume_enabled, get_unique_name

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)

logger = logging.getLogger(__name__)


def main():
    """完整的训练工作流程示例。"""
    
    # 1. 初始化配置
    logger.info("=" * 60)
    logger.info("初始化训练任务")
    logger.info("=" * 60)
    
    # 从环境变量获取配置
    env_config = get_env_config()
    project_name = env_config['PROJECT_NAME']
    run_name = env_config['RUN_NAME']
    dataset_uuid = env_config['DATASET_UUID']
    output_dir = env_config['OUTPUT_DIR']
    is_resume = is_resume_enabled()
    resume_from = env_config['RESUME_FROM']
    
    unique_name = get_unique_name(project_name, run_name)
    outdir_path = f"{output_dir}/{project_name}/{run_name}"
    
    logger.info(f"项目名称: {project_name}")
    logger.info(f"运行名称: {run_name}")
    logger.info(f"唯一标识: {unique_name}")
    logger.info(f"数据集UUID: {dataset_uuid}")
    logger.info(f"输出目录: {outdir_path}")
    logger.info(f"是否恢复: {is_resume}")
    
    # 2. 创建客户端
    client = RoboDockClient()
    
    try:
        # 3. 更新任务状态为处理中
        logger.info("更新任务状态为 processing...")
        client.update_run_status(
            unique_name=unique_name,
            status="processing",
            start_time=int(time.time() * 1000),
            finish_time=None
        )
        
        # 4. 下载数据集
        logger.info("=" * 60)
        logger.info("下载数据集")
        logger.info("=" * 60)
        
        dataset_info = client.get_dataset(dataset_uuid).get("data", {})
        if len(dataset_info) != 1:
            raise Exception(f"数据集信息无效: {dataset_info}")
        
        data_uuids = dataset_info[0]['linked_singledata_list']
        if not data_uuids:
            raise Exception("数据集中没有关联的数据")
        
        if len(data_uuids) > 1:
            logger.warning("发现多个关联数据，使用第一个")
        
        singledata_uuid = data_uuids[0]
        singledata_info = client.get_data(singledata_uuid).get("data", {})
        data_name = singledata_info.get("data_path", "").split("/")[-1]
        
        local_data_path = client.download_dataset(data_name, singledata_uuid, dirname="singledata")
        logger.info(f"数据集已下载到: {local_data_path}")
        
        # 5. 如果启用恢复，下载检查点
        if is_resume:
            logger.info("=" * 60)
            logger.info("下载检查点（恢复训练）")
            logger.info("=" * 60)
            
            run_info = client.get_run(resume_from).get("data", {})
            check_points = run_info.get("check_points_path", [])
            
            if not check_points or not isinstance(check_points, list):
                raise Exception("未找到可用的检查点")
            
            last_checkpoint = check_points[-1]
            logger.info(f"从检查点恢复: {last_checkpoint}")
            
            checkpoint_path = client.download_checkpoint(
                online_path=last_checkpoint,
                local_path=outdir_path
            )
            logger.info(f"检查点已下载到: {checkpoint_path}")
        
        # 6. 执行训练（这里用模拟代替实际训练）
        logger.info("=" * 60)
        logger.info("开始训练")
        logger.info("=" * 60)
        
        logger.info(f"数据路径: {local_data_path}")
        logger.info(f"输出路径: {outdir_path}")
        logger.info(f"训练步数: {env_config['NUM_TRAIN_STEPS']}")
        logger.info(f"批次大小: {env_config['BATCH_SIZE']}")
        
        # 这里应该调用实际的训练代码
        # 例如: subprocess.run(["python", "train.py", ...], check=True)
        
        # 模拟训练过程
        Path(outdir_path).mkdir(parents=True, exist_ok=True)
        with open(f"{outdir_path}/model.pt", "w") as f:
            f.write("模拟模型文件")
        
        logger.info("训练完成")
        
        # 7. 上传模型(同步/异步)
        logger.info("=" * 60)
        logger.info("上传模型")
        logger.info("=" * 60)
        #同步
        s3_key = client.upload_model(
            outdir_path=outdir_path,
            upload_path=f"models/{project_name}/{run_name}"
        )
        #异步
        # s3_key = client.upload_model_async(
        #     outdir_path=outdir_path,
        #     upload_path=f"models/{project_name}/{run_name}"
        # )
        logger.info(f"模型已上传: {s3_key}")
        
        # 8. 更新任务状态为成功
        logger.info("更新任务状态为 success...")
        client.update_run_status(
            unique_name=unique_name,
            status="success",
            finish_time=int(time.time() * 1000)
        )
        
        # 9. 清理本地文件
        logger.info("=" * 60)
        logger.info("清理本地文件")
        logger.info("=" * 60)
        
        if Path(outdir_path).exists():
            shutil.rmtree(outdir_path)
            logger.info(f"已删除本地检查点目录: {outdir_path}")
        
        logger.info("=" * 60)
        logger.info("训练任务完成！")
        logger.info("=" * 60)
        
    except Exception as e:
        error_info = f"训练过程出错: {e}"
        logger.error(error_info)
        
        # 更新任务状态为失败
        try:
            client.update_run_status(
                unique_name=unique_name,
                status="failed",
                fail_info=error_info,
                finish_time=int(time.time() * 1000)
            )
        except Exception as update_error:
            logger.error(f"更新失败状态出错: {update_error}")
        
        raise e


if __name__ == "__main__":
    main()
