"""
任务初始化工具
快速创建新任务的目录结构和配置文件模板
"""
import os
import json
import sys
import argparse


TASK_JSON_TEMPLATE = {
    "id": "",
    "category": "basic_processing",  # basic_processing, peak_analysis, phase_identification, compound_tasks
    "difficulty": "easy",  # easy, medium, hard
    
    "instruction": "请填写任务指令",
    
    "input": {
        "source_file": "../../data/source/DEMO01.MDI",
        "inject_to": "C:\\Users\\lzy\\Desktop\\DEMO01.MDI"
    },
    
    "output": {
        "expected_file": "result.txt",
        "collect_from": "C:\\Users\\lzy\\Desktop\\result.txt"
    },
    
    "evaluation": {
        "method": "xrd_data_compare",
        "ground_truth": "ground_truth/result.txt",
        "target_output": "agent_output/result.txt",
        "tolerance": 1e-4
    }
}


def init_task(task_id, project_root=".", force=False, category=None, difficulty=None, instruction=None):
    """
    初始化新任务
    
    Args:
        task_id: 任务ID
        project_root: 项目根目录
        force: 是否覆盖已存在的任务
        category: 任务类别（可选）
        difficulty: 任务难度（可选）
        instruction: 任务指令（可选）
    """
    task_dir = os.path.join(project_root, "tasks", task_id)
    
    # 检查是否已存在
    if os.path.exists(task_dir) and not force:
        print(f"❌ 任务目录已存在: {task_dir}")
        print("   使用 --force 参数强制覆盖")
        return False
    
    print(f"创建任务: {task_id}")
    print("=" * 60)
    
    # 创建目录结构
    directories = [
        task_dir,
        os.path.join(task_dir, "ground_truth"),
        os.path.join(task_dir, "human_demo"),
        os.path.join(task_dir, "human_demo", "screens"),
        os.path.join(task_dir, "agent_output")
    ]
    
    for directory in directories:
        os.makedirs(directory, exist_ok=True)
        print(f"✅ 创建目录: {os.path.relpath(directory, project_root)}")
    
    # 创建task.json
    task_config = TASK_JSON_TEMPLATE.copy()
    task_config["id"] = task_id
    
    # 更新可选参数
    if category:
        task_config["category"] = category
    if difficulty:
        task_config["difficulty"] = difficulty
    if instruction:
        task_config["instruction"] = instruction
    
    task_json_path = os.path.join(task_dir, "task.json")
    with open(task_json_path, 'w', encoding='utf-8') as f:
        json.dump(task_config, f, indent=2, ensure_ascii=False)
    
    print(f"✅ 创建配置: {os.path.relpath(task_json_path, project_root)}")
    
    # 创建README
    readme_content = f"""# 任务: {task_id}

## 任务信息
- **ID**: {task_id}
- **类别**: {task_config['category']}
- **难度**: {task_config['difficulty']}

## 指令
{task_config['instruction']}

## 数据采集状态
- [ ] 环境重置与文件注入
- [ ] 操作轨迹录制
- [ ] 结果文件收集
- [ ] 坐标转换处理
- [ ] 评测验证

## 采集命令
```bash
# 完整流程
python scripts/collect_task.py {task_id} --mode full

# 分步执行
python scripts/collect_task.py {task_id} --mode reset
python scripts/collect_task.py {task_id} --mode record
python scripts/collect_task.py {task_id} --mode collect
python scripts/process_trajectory.py {task_id}
python scripts/run_eval.py {task_id}
```

## 文件结构
```
{task_id}/
├── task.json              # 任务配置
├── ground_truth/          # 标准答案输出
├── human_demo/            # 人类操作轨迹
│   ├── actions_raw.json   # 原始轨迹（未转换坐标）
│   ├── actions.json       # 处理后轨迹（已转换坐标）
│   └── screens/           # 截图序列
└── agent_output/          # Agent输出（评测时使用）
```
"""
    
    readme_path = os.path.join(task_dir, "README.md")
    with open(readme_path, 'w', encoding='utf-8') as f:
        f.write(readme_content)
    
    print(f"✅ 创建说明: {os.path.relpath(readme_path, project_root)}")
    
    print("=" * 60)
    print("✅ 任务初始化完成！")
    print("\n📝 下一步:")
    print(f"   1. 编辑任务配置: {task_json_path}")
    print(f"   2. 确保输入文件存在:例如 {task_config['input']['source_file']}")
    print(f"   3. 开始数据采集: python scripts/tools/collect_task.py {task_id}")
    print("=" * 60 + "\n")
    
    return True


def main():
    parser = argparse.ArgumentParser(
        description="初始化新任务",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
使用示例:
  # 创建新任务
  python scripts/init_task.py smoothing_001
  
  # 强制覆盖已存在的任务
  python scripts/init_task.py smoothing_001 --force
        """
    )
    
    parser.add_argument("task_id", help="任务ID（建议格式: category_序号）")
    parser.add_argument("--project-root", default=".", help="项目根目录")
    parser.add_argument("--force", action="store_true", help="强制覆盖已存在的任务")
    
    args = parser.parse_args()
    
    success = init_task(args.task_id, args.project_root, args.force)
    sys.exit(0 if success else 1)


if __name__ == "__main__":
    main()