import yaml
from ultralytics import YOLO
import os
import glob


ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
def find_latest_run_dir(project_path='runs/detect/default_project'):
    """
    在指定的项目路径下，根据文件夹的修改时间找到最新的 'train' 目录。
    """
    if not os.path.exists(project_path):
        return None
    train_dirs = [d for d in glob.glob(os.path.join(project_path, 'train*')) if os.path.isdir(d)]
    if not train_dirs:
        return None
    latest_dir = max(train_dirs, key=os.path.getmtime)
    return latest_dir

def run_pipeline(config_path='config.yaml'):
    """
    读取配置文件并执行YOLO训练和导出流程。
    """
    # 1. --- 读取配置文件 ---
    try:
        with open(config_path, 'r', encoding='utf-8') as f:
            config = yaml.safe_load(f)
        print("✅ 配置文件加载成功！")
        print(f"项目名称: {config['project_name']}")
    except Exception as e:
        print(f"❌ 加载配置文件时发生错误: {e}")
        return

    # 提取核心配置
    base_model = config['base_model']
    data_yaml_path = config['data_yaml']
    project_name = config['project_name']
    
    model = YOLO(base_model)
    print(f"✅ 模型 '{base_model}' 初始化成功。")
    
    best_model_path = None

    # 2. --- 执行训练 ---
    if config.get('run_training', False):
        print("\n🚀 开始YOLO训练...")
        try:
            results = model.train(
                data=data_yaml_path,
                epochs=config['epochs'],
                imgsz=config['imgsz'],
                batch=config['batch_size'],
                workers=config['workers'],
                project=project_name,
                name='train'
            )
            
            final_results = results[0] if isinstance(results, (list, tuple)) else results
            
            assert final_results is not None and hasattr(final_results, 'save_dir'), \
                "训练未产生一个包含保存目录的有效结果对象。"
            
            best_model_path = os.path.join(final_results.save_dir, 'weights/best.pt')

            print(f"✅ 训练完成！最佳模型已保存在: {best_model_path}")
        except Exception as e:
            print(f"❌ 训练过程中发生错误: {e}")
            return
    else:
        print("\n⏩ 根据配置，跳过训练步骤。")
        if config.get('run_export', False):
            print("   正在查找最新的训练结果...")
            project_path = os.path.join(ROOT_DIR, project_name)
            latest_run_dir = find_latest_run_dir(project_path)
            
            if latest_run_dir:
                potential_model_path = os.path.join(latest_run_dir, 'weights', 'best.pt')
                if os.path.exists(potential_model_path):
                    best_model_path = potential_model_path
                    print(f"✅ 成功找到最新的模型: {best_model_path}")
                else:
                    print(f"❌ 在最新的训练目录 '{latest_run_dir}' 中未找到 'best.pt' 文件。")
            else:
                print(f"❌ 在项目 '{project_name}' 中未找到任何过往的训练结果。")

    # 3. --- 执行导出 ---
    print(f"root_dir:{ROOT_DIR}")
    # print(f"run_export:{config.get('run_export', False)}")
    print(f"best model path:{best_model_path}")
    if config.get('run_export', False) and best_model_path:
        
        # --- 修改开始：增加最终的安全检查 ---
        # 确保 best_model_path 是一个字符串，而不是元组或列表
        if isinstance(best_model_path, (list, tuple)):
            print(f"⚠️ 检测到模型路径为序列类型，自动提取第一个元素。原始值: {best_model_path}")
            best_model_path = best_model_path[0]
        # --- 修改结束 ---

        print(f"\n🚀 开始将模型 '{best_model_path}' 导出为 {config['export_format']} 格式...")
        try:
            model_to_export = YOLO(best_model_path)
            
            model_to_export.export(
                format=config['export_format'],
                imgsz=config['imgsz'],
                half=config.get('half_precision', False)
            )
            
            exported_file_name = os.path.basename(best_model_path).replace('.pt', f".{config['export_format']}")
            exported_file_path = os.path.join(os.path.dirname(best_model_path), exported_file_name)
            print(f"✅ 导出成功！文件已保存在: {exported_file_path}")
        except Exception as e:
            print(f"❌ 导出过程中发生错误: {e}")
    elif config.get('run_export', False):
         print("\n⏩ 跳过导出步骤，因为未找到有效的模型路径。")
    else:
        print("\n⏩ 根据配置，跳过导出步骤。")
        
    print("\n🎉 自动化流程执行完毕！")


if __name__ == '__main__':
    run_pipeline()