import argparse
import glob
import os
import multiprocessing as mp
from pathlib import Path
import re

from until.collect_import_modeling import expand_modeling_imports,remove_imports_and_rewrite,save_results_to_txt
from until.rewrite_child_classes import rewrite_child_classes
from until.rename_identifiers import rename_identifiers

AUTO_GENERATED_MESSAGE = """#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
#           This file was automatically generated from {relative_path}.
#               Do NOT edit this file manually as any edits will be overwritten by the generation of
#             the file from the modular. If any change should be done, please apply the change to the
#                          {short_name} file directly. One of our CI enforces this.
#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
"""

# --- 核心转换逻辑封装 ---
def run_converter(file_to_parse_str: str):
    """
    对单个 modular 文件执行完整的转换流程。
    这个函数是并行处理的基本单元。
    """
    print(file_to_parse_str)
    file_to_parse = Path(file_to_parse_str)
    
    # --- 动态确定模型名称和输出路径 ---
    # 假设：目标模型名可以从文件名推断，例如 "modular_qwen2.py" -> "qwen2"
    to_name = file_to_parse.stem.replace("modular_", "")
    # 假设：源模型名在模块化文件中有定义或约定俗成（这里我们先硬编码为 "llama" 作为示例）
    # 一个更健壮的实现会从 file_to_parse 文件内容中解析出 from_name
    
    # 输出文件将与输入文件在同一目录下，例如 "modeling_qwen2.py"
    output_file = file_to_parse.parent / f"modeling_{to_name}.py"
    temp_merged_file = file_to_parse.parent / (output_file.stem + "_temp_merged.py")

    print(f"--- 开始转换: {to_name} ---")
    print(f"  输入文件: '{file_to_parse}'")
    print(f"  最终输出: '{output_file}'")

    # 步骤 1: 收集并展开 import
    expanded_code , dependencies,from_name= expand_modeling_imports(file_to_parse)
    #save_results_to_txt(dependencies, "modeling_imports_dependencies.txt")
    #save_results_to_txt(expanded_code, "modeling_imports_results.txt")
    #print(from_name)
    # 步骤 2: 重写子类并生成中间文件
    relative_path = re.search(
                    r"(transformers/.*|examples/.*)", os.path.abspath(file_to_parse).replace("\\", "/")
                ).group(1)
    formatted_message=AUTO_GENERATED_MESSAGE.format(relative_path=relative_path, short_name=os.path.basename(relative_path))
    rewrite_child_classes(expanded_code, file_to_parse,formatted_message, temp_merged_file,rename_map={
            "llama": "qwen2"  # 只需要提供小写形式！
        })
    remove_imports_and_rewrite(temp_merged_file)
    # 步骤 3: 全局重命名
    try:
        merged_code = temp_merged_file.read_text(encoding="utf-8")
        final_code = rename_identifiers(merged_code, from_name, to_name)
        output_file.write_text(final_code, encoding="utf-8")
        print(f"  ✅ 转换成功，最终代码已写入 '{output_file}'。")
    except FileNotFoundError:
        print(f"  ❌ [错误] 找不到中间文件 '{temp_merged_file}'，无法进行重命名。")
    finally:
        # 清理临时文件
        if temp_merged_file.exists():
            temp_merged_file.unlink()
    
    print(f"--- 转换结束: {to_name} ---\n")


# --- 主执行逻辑 ---
def main():
    parser = argparse.ArgumentParser(
        description="将模块化的模型定义文件（modular_*.py）转换为独立的模型文件（modeling_*.py）。"
    )
    parser.add_argument(
        "files",
        nargs="*",
        help="要转换的模块化文件列表（可选的位置参数）。",
    )
    parser.add_argument(
        "--files-to-parse", "-f",
        dest="files_to_parse", # 明确指定存储的目的地
        default=[], # 默认值改为空列表
        nargs="+",
        help="要转换的模块化文件列表。可使用 'all' 或 'examples' 关键字。",
    )
    parser.add_argument(
        "--num_workers", "-w",
        default=-1,
        type=int,
        help="使用的进程数。默认为 -1，代表使用所有 CPU核心。",
    )
    args = parser.parse_args()

    # 合并位置参数和可选参数，以可选参数优先
    files_to_parse = args.files_to_parse if args.files_to_parse else args.files
    if not files_to_parse:
         files_to_parse = ["all"] # 如果未提供任何文件，则默认为 'all'

    num_workers = mp.cpu_count() if args.num_workers == -1 else args.num_workers

    # --- 解析文件路径 ---
    print(">>> 正在解析需要转换的文件...")
    if files_to_parse == ["all"]:
        from pathlib import Path
        # 确定项目根目录
        SCRIPT_DIR = Path(__file__).resolve().parent
        PROJECT_ROOT = SCRIPT_DIR.parent.parent.parent
        # 使用绝对路径进行搜索
        search_path = PROJECT_ROOT / "paddleformers/transformers/**/modular_*.py"
        files_to_parse = glob.glob(str(search_path), recursive=True)
    elif files_to_parse == ["examples"]:
        # 查找所有 examples 目录下的 modular 文件
        files_to_parse = glob.glob("examples/**/modular_*.py", recursive=True)
    else:
        # 将模型简称（如 qwen2）解析为完整路径
        resolved_files = []
        for model_name in files_to_parse:
            if not os.path.exists(model_name):
                # 尝试在 models 目录下构建路径
                full_path = os.path.join("src", "transformers", "models", model_name, f"modular_{model_name}.py")
                if not os.path.isfile(full_path):
                     # 如果找不到，尝试在 examples 目录下构建
                    full_path = os.path.join("examples", "modular-transformers", f"modular_{model_name}.py")
                
                if not os.path.isfile(full_path):
                    raise ValueError(f"无法为 '{model_name}' 找到模块化文件。请提供完整路径或确认文件名正确。")
                resolved_files.append(full_path)
            else:
                resolved_files.append(model_name)
        files_to_parse = resolved_files
    
    if not files_to_parse:
        print("未找到任何需要转换的文件。")
        return
        
    print(f"发现 {len(files_to_parse)} 个文件待处理。")

    
    ordered_files= [files_to_parse]
    print(ordered_files)

    # --- 按依赖顺序并行处理 ---
    for i, dependency_level_files in enumerate(ordered_files):
        print(f"\n>>> 开始处理依赖层级 {i+1}/{len(ordered_files)} ({len(dependency_level_files)} 个文件)...")
        workers = min(num_workers, len(dependency_level_files))
        if workers > 0:
            with mp.Pool(workers) as pool:
                pool.map(run_converter, dependency_level_files)
    
    print("\n--- 所有转换任务已完成 ---")


if __name__ == "__main__":
    main()