import pandas as pd
import os
import re


def extract_cif_from_xlsx(
        xlsx_path: str,
        output_dir: str,
        naming_mode: str = 'formula',
        name_col: int = 0,
        cif_col: int = 1,
        prefix: str = 'wjy'
):
    """
    从 XLSX 文件中提取 CIF 数据并保存为单独的 .cif 文件。

    Args:
        xlsx_path (str): 输入的 XLSX 文件的路径。
        output_dir (str): 输出 .cif 文件的文件夹路径。
        naming_mode (str, optional): CIF 文件的命名模式。
                                     可选值为 'formula' (使用第一列的名字) 或
                                     'auto' (使用前缀+自动递增编号)。
                                     默认为 'formula'。
        name_col (int, optional): 包含文件名的列的索引（从0开始）。默认为 0。
        cif_col (int, optional): 包含 CIF 内容的列的索引（从0开始）。默认为 1。
        prefix (str, optional): 在 'auto' 命名模式下使用的文件名前缀。默认为 'wjy'。

    Raises:
        FileNotFoundError: 如果指定的 xlsx_path 文件不存在。
        ValueError: 如果指定的 naming_mode 无效。
        Exception: 处理过程中发生的其他错误。
    """
    # --- 1. 参数校验和准备 ---
    if not os.path.exists(xlsx_path):
        raise FileNotFoundError(f"错误: 输入文件未找到 -> {xlsx_path}")

    if naming_mode not in ['formula', 'auto']:
        raise ValueError(f"错误: 'naming_mode' 参数必须是 'formula' 或 'auto'，但收到了 '{naming_mode}'")

    # 创建输出目录（如果不存在）
    os.makedirs(output_dir, exist_ok=True)
    print(f"CIF 文件将保存到: {output_dir}")

    try:
        # --- 2. 读取 XLSX 文件 ---
        # header=None 表示第一行不是标题，将其作为数据读取
        df = pd.read_excel(xlsx_path, header=None)

        # 跳过原始文件的表头行（'formula', 'cif'）
        if str(df.iloc[0, name_col]).strip().lower() == 'formula' and str(df.iloc[0, cif_col]).strip().lower() == 'cif':
            df = df.iloc[1:]
            print("检测到并跳过了表头行。")

        # --- 3. 遍历数据并生成文件 ---
        success_count = 0
        for index, row in df.iterrows():
            # 获取文件名和 CIF 内容
            formula_name = str(row[name_col])
            cif_content = str(row[cif_col])

            # 跳过内容为空的行
            if pd.isna(row[name_col]) or pd.isna(row[cif_col]) or not cif_content.strip():
                print(f"警告: 第 {index + 2} 行数据不完整，已跳过。")
                continue

            # --- 4. 根据命名模式确定文件名 ---
            if naming_mode == 'formula':
                # 清理文件名，替换掉不适合做文件名的特殊字符
                # 例如：将 (PO4)3 替换为 _PO4_3，将 / 替换为 _
                safe_filename = re.sub(r'[\\/*?:"<>|()]', '_', formula_name)
                filename = f"{safe_filename}.cif"
            else:  # naming_mode == 'auto'
                # 使用 format 方法来确保编号格式统一，例如 001, 002
                filename = f"{prefix}_{success_count + 1:03d}.cif"

            # 构造完整的输出文件路径
            output_path = os.path.join(output_dir, filename)

            # --- 5. 写入 CIF 文件 ---
            try:
                with open(output_path, 'w', encoding='utf-8') as f:
                    f.write(cif_content)
                success_count += 1
            except IOError as e:
                print(f"错误: 无法写入文件 {output_path}。原因: {e}")

        print(f"\n处理完成！成功提取并生成了 {success_count} 个 CIF 文件。")

    except Exception as e:
        print(f"处理 XLSX 文件时发生错误: {e}")


# --- 函数使用示例 ---
if __name__ == '__main__':
    # 假设您的 XLSX 文件名为 'materials.xlsx'，且与此脚本在同一目录下
    source_xlsx_file = 'input/cif_dataset.xlsx'

    # 检查示例文件是否存在，如果不存在则创建一个
    if not os.path.exists(source_xlsx_file):
        print(f"未找到示例文件 '{source_xlsx_file}'，正在创建一个...")
        example_data = {
            'formula': ['Li3Al0.3Ti1.7(PO4)3', 'Li6.5La3Zr1.75W0.25O12', 'Invalid/Name*Test'],
            'cif': ['# CIF Data for Li3Al0.3...\n_atom_site_type_symbol\n Li\n Al\n Ti\n P\n O',
                    '# CIF Data for Li6.5La3...\n_symmetry_space_group_name_H-M \'I a -3 d\'',
                    '# CIF Data for Invalid Name Test']
        }
        pd.DataFrame(example_data).to_excel(source_xlsx_file, index=False, header=True)
        print("示例文件创建成功。")

    # --- 示例 1: 使用第一列的 'formula' 命名 ---
    # print("\n--- 示例 1: 使用 'formula' 命名模式 ---")
    # output_folder_1 = 'cif_by_formula'
    # extract_cif_from_xlsx(
    #     xlsx_path=source_xlsx_file,
    #     output_dir=output_folder_1,
    #     naming_mode='formula'
    # )

    # --- 示例 2: 使用 'wjy+编号' 自动命名 ---
    print("\n--- 示例 2: 使用 'auto' 命名模式 ---")
    output_folder_2 = 'cif_by_auto'
    extract_cif_from_xlsx(
        xlsx_path=source_xlsx_file,
        output_dir=output_folder_2,
        naming_mode='auto',
        prefix='wjy'
    )