import os
import json
import glob

def main():
    base_dir = r"c:\Users\11304\gitlab\cybertogether\Agent\examples\process_research\output"
    
    # 遍历 output 目录下的所有子文件夹
    if not os.path.exists(base_dir):
        print(f"Directory not found: {base_dir}")
        return

    fixed_count = 0
    skip_count = 0

    for item in os.listdir(base_dir):
        sub_dir = os.path.join(base_dir, item)
        if not os.path.isdir(sub_dir):
            continue
            
        strategy_path = os.path.join(sub_dir, "strategy.json")
        if not os.path.exists(strategy_path):
            continue
            
        # 读取 strategy.json
        try:
            with open(strategy_path, "r", encoding="utf-8") as f:
                strategy_data = json.load(f)
        except Exception as e:
            print(f"Error reading {strategy_path}: {e}")
            continue

        # 如果已经有了，跳过
        if "requirement" in strategy_data:
            print(f"[{item}] skip: already has requirement")
            skip_count += 1
            continue
            
        # 寻找对应的 case.json 提取 requirement
        requirement_text = None
        
        # 1. 尝试直接找 case.json
        case_path = os.path.join(sub_dir, "case.json")
        if os.path.exists(case_path):
            try:
                with open(case_path, "r", encoding="utf-8") as f:
                    case_data = json.load(f)
                    requirement_text = case_data.get("requirement")
            except Exception:
                pass
                
        # 2. 如果没找到或 case.json 不存在，并且存在分散的 case_*.json
        if not requirement_text:
            case_files = glob.glob(os.path.join(sub_dir, "case_*.json"))
            for cf in case_files:
                try:
                    with open(cf, "r", encoding="utf-8") as f:
                        cf_data = json.load(f)
                        if "requirement" in cf_data:
                            requirement_text = cf_data["requirement"]
                            break
                except Exception:
                    continue
                    
        if not requirement_text:
            print(f"[{item}] ERROR: could not find 'requirement' in any case files")
            continue
            
        # 更新 strategy.json
        # 为了保证 requirement 在最开头，创建一个新的 dict
        new_strategy_data = {"requirement": requirement_text}
        new_strategy_data.update(strategy_data)
        
        try:
            with open(strategy_path, "w", encoding="utf-8") as f:
                json.dump(new_strategy_data, f, ensure_ascii=False, indent=4)
            print(f"[{item}] fixed: added requirement successfully")
            fixed_count += 1
        except Exception as e:
            print(f"[{item}] ERROR writing back strategy data: {e}")

    print(f"\nSummary: Fixed {fixed_count} files, Skipped {skip_count} files.")

if __name__ == "__main__":
    main()