import os


def create_sh_recursive(base_folder, tool_path="tool", relative_depth=2):
    """
    递归遍历文件夹,为每个包含.cif文件的文件夹生成analyze.sh脚本,
    并在基础文件夹下创建一个sh_all.sh来执行所有脚本。
    """
    # 用于收集所有生成的analyze.sh脚本的相对路径
    analyze_sh_paths = []

    def process_folder(folder_path, current_depth=0):
        print(f"处理文件夹: {folder_path}")

        # 获取当前文件夹名称
        folder_name = os.path.basename(folder_path)

        # 检查当前文件夹是否包含.cif文件
        # 注意：这里我们只关心当前层级下的cif文件
        cif_files = [f for f in os.listdir(folder_path) if
                     f.endswith('.cif') and os.path.isfile(os.path.join(folder_path, f))]
        has_cif_files = len(cif_files) > 0

        # 如果当前文件夹包含.cif文件,生成脚本
        if has_cif_files:
            # 计算正确的工具路径(根据深度增加../)
            dots = "../" * (relative_depth + current_depth)
            tool_relative_path = f"{dots}{tool_path}"

            # --- 修改开始: 修正Anion识别逻辑 ---
            # 如果结构是 Anion/MaterialID/file.cif，此时folder_name是MaterialID
            # 我们需要上一级目录的名字作为Anion (例如 'O', 'S', 'Cl') 来寻找对应的 .yaml 文件
            # 简单的判断逻辑：如果当前文件夹名字主要由数字组成(或者是ID格式)，且包含cif文件，我们假设其父目录是Anion类型
            # 或者更直接的逻辑：在你的新结构中，包含cif的文件夹必定是底层文件夹，其父目录必定是Anion
            parent_dir_name = os.path.basename(os.path.dirname(folder_path))

            # 这里做一个简单的保护，如果是在第一层(比如直接在O文件夹下有cif)，保持原状，否则使用父目录名
            # 在新结构下，cif总是在 '.../O/141/141.cif'，所以anion应该是 parent_dir_name ('O')
            if parent_dir_name in ['O', 'S', 'Cl', 'Br'] or folder_name not in ['O', 'S', 'Cl', 'Br']:
                anion = parent_dir_name
            else:
                anion = folder_name
            # --- 修改结束 ---

            # 生成脚本文件路径
            sh_file_path = os.path.join(folder_path, "analyze.sh")

            # 创建脚本
            with open(sh_file_path, 'w') as sh_file:
                sh_file.write('#!/bin/bash\n')
                for filename in cif_files:
                    # --- 修改开始: 输出重定向到 log.txt ---
                    command = f"python {tool_relative_path}/analyze_voronoi_nodes.py {filename} -i {tool_relative_path}/{anion}.yaml > log.txt\n"
                    # --- 修改结束 ---
                    sh_file.write(command)

            # 将此脚本添加到收集器中
            rel_path = os.path.relpath(folder_path, base_folder)
            analyze_sh_paths.append(rel_path)

            print(f"生成脚本: {sh_file_path} (工具路径: {tool_relative_path}, Anion: {anion})")

        # 获取子文件夹列表
        subdirs = [d for d in os.listdir(folder_path) if os.path.isdir(os.path.join(folder_path, d))]

        # 处理复合阴离子文件夹的特殊情况
        if "+" in folder_name:
            elements = folder_name.split("+")
            for element in elements:
                element_dir = os.path.join(folder_path, element)
                # 如果对应元素的子文件夹不存在,创建它
                if not os.path.exists(element_dir):
                    os.makedirs(element_dir)
                    print(f"创建子文件夹: {element_dir}")
                # 确保这个子文件夹被包含在递归处理列表中
                if element not in subdirs:
                    subdirs.append(element)

        # 递归处理所有子文件夹
        for subdir in subdirs:
            subdir_path = os.path.join(folder_path, subdir)
            process_folder(subdir_path, current_depth + 1)

    # 开始递归处理
    process_folder(base_folder)

    # 创建sh_all.sh脚本
    sh_all_path = os.path.join(base_folder, "sh_all.sh")
    with open(sh_all_path, 'w') as sh_all:
        sh_all.write('#!/bin/bash\n\n')
        sh_all.write(f'# process all analyze.sh\n\n')

        # 记录初始目录
        sh_all.write('# remember current dir\n')
        sh_all.write('INITIAL_DIR=$(pwd)\n\n')

        # 为每个包含analyze.sh的目录添加命令
        for path in analyze_sh_paths:
            sh_all.write(f'echo "process {path}/analyze.sh"\n')
            sh_all.write(f'cd "{path}"\n')
            sh_all.write('bash analyze.sh\n')
            sh_all.write(f'cd "$INITIAL_DIR"\n\n')

        # 添加完成消息
        sh_all.write('echo "done!"\n')

    # 修改权限使脚本可执行
    os.chmod(sh_all_path, 0o755)
    print(f"生成总执行脚本: {sh_all_path}")
    print("所有脚本生成完成!")


if __name__ == '__main__':
    create_sh_recursive("../data/after_step1")