import os
from pathlib import Path
import sys
import inspect
from typing import Union, Tuple, Optional
import shutil
import subprocess
import zipfile
from PIL import Image


def code2prompt_cmd():
    """
    交互式生成code2prompt命令，优化了排除文件类型的输入方式，更友好的交互体验
    """
    print("===== code2prompt命令生成工具 =====")

    # 1. 获取项目路径（待处理的代码目录）
    while True:
        project_path = input("请输入项目根目录路径：").strip()
        if os.path.isdir(project_path):
            break
        print(f"错误：路径 '{project_path}' 不存在或不是目录，请重新输入！")

    # 2. 获取输出路径（生成文件的保存目录）
    while True:
        output_dir = input("请输入输出文件保存目录：").strip()
        if os.path.isdir(output_dir):
            break
        # 询问是否创建不存在的目录
        create = input(f"目录 '{output_dir}' 不存在，是否创建？(y/n)：").strip().lower()
        if create == 'y':
            os.makedirs(output_dir, exist_ok=True)
            print(f"已创建目录：{output_dir}")
            break
        print("请重新输入输出目录！")

    # 3. 获取输出文件名
    while True:
        filename = input("请输入输出文件名（例如：ai_prompt.md）：").strip()
        if filename:
            # 确保文件名包含扩展名（默认.md）
            if '.' not in filename:
                filename += '.md'
                print(f"自动补充扩展名，文件名变为：{filename}")
            break
        print("文件名不能为空，请重新输入！")

    # 4. 是否需要显示行号
    line_number = input("是否需要包含行号？(y/n，默认y)：").strip().lower()
    line_number_flag = "--line-number" if (not line_number or line_number == 'y') else ""

    # 5. 优化：排除文件类型（支持空格分隔，自动处理为*.xxx格式，默认选项更清晰）
    default_exclude_types = ["txt", "logs", "csv", "cmd", "sh", "md"]  # 仅保留类型名
    default_exclude_display = "、".join(default_exclude_types)  # 显示为“txt、logs、csv...”
    exclude_input = input(
        f"请输入需要排除的文件类型（用空格分隔；直接回车使用默认：{default_exclude_display}）："
    ).strip()

    if not exclude_input:
        # 使用默认排除类型
        exclude_files = ",".join([f"*.{t}" for t in default_exclude_types])
    else:
        # 处理用户输入（空格分隔转成*.xxx,*.yyy格式）
        exclude_types = exclude_input.split()
        exclude_files = ",".join([f"*.{t.lstrip('.')}" for t in exclude_types])  # 兼容用户输入带.的情况（如.txt）

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

    # 生成最终命令（处理路径中的空格，用双引号包裹）
    cmd_parts = [
        "code2prompt",
        f'-p "{project_path}"',
        f'-o "{output_path}"',
        line_number_flag,
        f'-e "{exclude_files}"'
    ]
    # 过滤空值（如不显示行号时）
    cmd = ' '.join(part for part in cmd_parts if part)

    print("\n===== 生成的命令如下 =====")
    print(cmd)
    return cmd
def add_path_to_sys(path: Union[str, Path]) -> Optional[str]:
    """
    将指定路径添加到sys.path中（自动转换为绝对路径，避免重复添加）
    :param path: 要添加的路径（支持字符串或Path对象）
    :return: 成功添加返回路径字符串，已存在返回None
    """
    # 转换为绝对路径字符串（统一格式，避免因相对路径导致的重复）
    abs_path = str(Path(path).resolve())

    # 检查是否已存在，避免重复添加
    if abs_path not in sys.path:
        sys.path.insert(0, abs_path)  # 插入到首位，优先搜索
        print(f"已添加路径到搜索路径：{abs_path}")
        return abs_path
    else:
        print(f"路径已在搜索路径中，无需重复添加：{abs_path}")
        return None
def remove_path_from_sys(path: Union[str, Path]) -> Tuple[bool, str]:
    """
    从sys.path中移除指定路径（自动转换为绝对路径，确保匹配）
    :param path: 要移除的路径（支持字符串或Path对象）
    :return: (是否成功移除, 处理后的绝对路径)
    """
    # 转换为绝对路径字符串（与添加时的格式保持一致）
    abs_path = str(Path(path).resolve())

    # 尝试移除路径
    if abs_path in sys.path:
        sys.path.remove(abs_path)
        print(f"已从搜索路径中移除：{abs_path}")
        return True, abs_path
    else:
        print(f"路径不在搜索路径中，无法移除：{abs_path}")
        return False, abs_path
def copy_from_dataset(target_path):
    ''''
    从固定源文件夹复制文件到目标路径
    '''
    source_folder = r"D:\时序预测\常用数据集"
    if not os.path.exists(source_folder):
        print(f"错误：源文件夹不存在 → {source_folder}")
        return
    if not os.path.exists(target_path):
        os.makedirs(target_path)
    if os.path.isdir(target_path):
        for filename in os.listdir(source_folder):
            source_file = os.path.join(source_folder, filename)
            if os.path.isfile(source_file):
                try:
                    shutil.copy2(source_file, target_path)
                    print(f"已复制 {filename} 到 {target_path}")
                except Exception as e:
                    print(f"复制 {filename} 失败：{e}")
    else:
        filename = os.path.basename(target_path)
        source_file = os.path.join(source_folder, filename)
        if os.path.isfile(source_file):
            try:
                shutil.copy2(source_file, target_path)
                print(f"已复制到：{target_path}")
            except Exception as e:
                print(f"复制失败：{e}")
        else:
            print(f"源文件夹中不存在文件：{filename}")

def pandoc_tool(src_file=None, target_format=None, out_dir=None):
    '''
    Pandoc格式转换工具,支持交互/传参双模式,转换tex自动植入中文+表格+边距配置,返回最终输出文件绝对路径
    :param src_file: 待转换文件完整路径,传参则静默执行,不传则弹窗输入
    :param target_format: 目标格式,支持.md .html .tex .docx,传参则静默执行,不传则弹窗选择
    :param out_dir: 输出目录完整路径,可选参数,不传则输出到源文件同级目录
    :return: 成功返回输出文件绝对路径,失败返回None
    '''
    param_mode = True if src_file and target_format else False
    choice = None
    if not param_mode:
        print("=" * 65)
        print("        Pandoc 格式转换工具【终极完整版】支持：独立文件夹+专属media+自动提图 ✅")
        print("        ✅ 中文配置+图片提取+路径适配+独立目录 全功能集成 ✅")
        print("=" * 65)
        print("\n✅ 支持：直接输文件名(同目录)｜完整路径｜路径带/不带双引号 都可以")
        src_file = input("📌 请输入需要转换的文件完整路径：")
        src_file = src_file.strip().strip('"').strip('“').strip('”')
        if not os.path.exists(src_file):
            print(f"\033[31m❌ 错误：找不到该文件，请检查路径是否输入正确 → {src_file}\033[0m")
            input("\n按任意键退出...")
            return None
        print("\n\n📌 请选择【要转换为】的目标格式 (输入数字回车即可)")
        print("1 → 转换成 Markdown(.md)  【最常用，给AI处理首选】")
        print("2 → 转换成 HTML(.html)    【保留完整排版/图片，网页可用】")
        print("3 → 转换成 LaTeX(.tex)    【学术公式专用，论文最优｜自动加中文配置+提图】")
        print("4 → 转换成 Word(.docx)    【反向转换，md/html/tex转回Word】")
        while True:
            choice = input("\n你的选择(1-4)：")
            if choice in ["1", "2", "3", "4"]: break
            print("❌ 输入错误！仅支持输入 1、2、3、4 其中一个数字")
        target_format = {"1": ".md", "2": ".html", "3": ".tex", "4": ".docx"}[choice]

    src_file = src_file.strip().strip('"').strip('“').strip('”')
    if not os.path.exists(src_file):
        if param_mode:
            print(f"转换失败：文件不存在 -> {src_file}")
        else:
            print(f"\033[31m❌ 错误：找不到该文件，请检查路径是否输入正确 → {src_file}\033[0m");
            input("\n按任意键退出...")
        return None

    # ============ 核心修改1：创建 输出目录/同名独立文件夹 的目录结构 ============
    src_dir, src_name = os.path.split(src_file)
    file_prefix = os.path.splitext(src_name)[0]
    # 根输出目录：传了out_dir用out_dir，没传就用源文件同级目录
    root_out_dir = out_dir if (out_dir and os.path.exists(out_dir)) else src_dir
    # 核心：为每个文件创建【同名独立文件夹】
    single_file_dir = os.path.join(root_out_dir, file_prefix)
    if not os.path.exists(single_file_dir):
        os.makedirs(single_file_dir)
    # 最终转换后的文件 放在这个独立文件夹里
    tar_file = os.path.join(single_file_dir, f"{file_prefix}{target_format}")
    pandoc_cmd = ['pandoc', '-s', src_file, '-o', tar_file]

    try:
        # 执行转换命令
        subprocess.run(pandoc_cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, encoding='utf-8')
        if param_mode:
            print(f"转换成功: {os.path.basename(src_file)} -> {target_format}")
        else:
            print(f"\033[32m✅ 转换成功！已创建独立文件夹存放文件 ✅\033[0m")

        # ============ 核心修改2：TEX专属配置 + 自动提取图片 + 图片路径适配 ============
        if target_format == '.tex':
            # 1. 插入原版完美中文+表格+边距配置
            insert_code = r"""
\usepackage{ctex}
\usepackage{fontspec}
\usepackage{upgreek}
\setCJKmainfont[AutoFakeBold=2.5]{SimSun}
\setmainfont{Times New Roman}
\usepackage{geometry}
\geometry{left=2cm,right=2cm,top=2cm,bottom=2cm}
\usepackage{booktabs}
\usepackage{array}
\usepackage{graphicx}
\renewcommand{\arraystretch}{1.2}"""
            with open(tar_file, 'r', encoding='utf-8') as f:
                all_lines = f.readlines()
            new_lines = []
            inserted = False
            for line in all_lines:
                new_lines.append(line)
                if '{article}' in line and not inserted:
                    new_lines.append(insert_code)
                    inserted = True
            with open(tar_file, 'w', encoding='utf-8') as f:
                f.writelines(new_lines)

            # 2. 调用提取图片函数：提取到 当前tex文件同级的media文件夹
            media_folder = os.path.join(single_file_dir, "media")
            extract_word_images(src_file, media_folder)

            # 3. 自动修正tex里的图片路径为相对路径 media/xxx，保证编译能找到图片
            with open(tar_file, 'r', encoding='utf-8') as f:
                tex_content = f.read()
            # 适配pandoc默认生成的media路径，统一修正为当前独立文件夹的media
            tex_content = tex_content.replace("media/", "./media/")
            with open(tar_file, 'w', encoding='utf-8') as f:
                f.write(tex_content)

            if param_mode:
                print(f"✅ 配置成功: 中文配置+图片提取+路径适配完成")
            else:
                print(f"\033[33m✅ ✅ ✅ 中文配置+图片提取+路径适配 全部完成！✅ ✅ ✅\033[0m")

        # ============ 输出文件路径信息 ============
        abs_tar_path = os.path.abspath(tar_file)
        if param_mode:
            print(f"文件路径: {abs_tar_path}")
        else:
            print(f"\033[32m✅ 最终文件位置：{abs_tar_path}\033[0m")
        return abs_tar_path

    except FileNotFoundError:
        if param_mode:
            print("转换失败：Pandoc环境变量未配置")
        else:
            print("\033[31m❌ 转换失败：Pandoc环境变量未配置好，请检查Path仅填文件夹路径！\033[0m")
    except subprocess.CalledProcessError:
        if param_mode:
            print("转换失败：文件格式不支持/文件损坏")
        else:
            print("\033[31m❌ 转换失败：文件格式不支持转换/文件已损坏！\033[0m")
    except Exception as e:
        if param_mode:
            print(f"转换失败：{str(e)}")
        else:
            print(f"\033[31m❌ 其他错误：{str(e)}\033[0m")

    if not param_mode: input("\n转换完成，按【任意键】退出工具...")
    return None
def extract_word_images(docx_path, output_folder="media"):
    """
    提取Word文档(.docx)中的所有图片，自动放进指定文件夹 + 自动将EMF格式批量转为PNG + 删除原EMF文件
    适配你的Pandoc转换工具，彻底解决LaTeX的EMF图片报错问题
    :param docx_path: Word文档路径
    :param output_folder: 输出文件夹，默认是media
    :return: None
    """
    if not docx_path.lower().endswith(".docx"):
        return
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    with zipfile.ZipFile(docx_path, 'r') as zf:
        for file_path in zf.namelist():
            if file_path.startswith("word/media/") and len(file_path) > len("word/media/"):
                img_name = os.path.basename(file_path)
                zf.extract(file_path, output_folder)
                src_img = os.path.join(output_folder, file_path)
                dst_img = os.path.join(output_folder, img_name)
                if os.path.exists(src_img):
                    shutil.move(src_img, dst_img)
    temp_word_dir = os.path.join(output_folder, "word")
    if os.path.exists(temp_word_dir):
        shutil.rmtree(temp_word_dir, ignore_errors=True)

    # 批量转换EMF->PNG+删除原文件
    emf_files = []
    for img_file in os.listdir(output_folder):
        img_file_path = os.path.join(output_folder, img_file)
        if os.path.isfile(img_file_path) and img_file.lower().endswith(".emf"):
            try:
                img_name_no_suffix = os.path.splitext(img_file)[0]
                png_save_path = os.path.join(output_folder, f"{img_name_no_suffix}.png")
                img = Image.open(img_file_path)
                img.save(png_save_path, format="PNG")
                emf_files.append(img_file_path)
                print(f"🔄 自动转换：{img_file} -> {img_name_no_suffix}.png")
            except Exception as e:
                print(f"⚠️ 转换失败 {img_file}：{str(e)}")
    for emf_file in emf_files:
        if os.path.exists(emf_file):
            os.remove(emf_file)
            print(f"✅ 删除原EMF：{os.path.basename(emf_file)}")
    print(f"✅ 图片提取+转换完成：{os.path.abspath(output_folder)}")

def get_model_size_mb(model):
    # 计算模型参数在内存中所占的实际大小(MB)
    param_size = 0
    for param in model.parameters():
        param_size += param.nelement() * param.element_size()
    buffer_size = 0
    for buffer in model.buffers():
        buffer_size += buffer.nelement() * buffer.element_size()
    size_all_mb = (param_size + buffer_size) / 1024**2
    print(f'Model size: {size_all_mb:.3f} MB')
    return size_all_mb

if __name__ == "__main__":
    pandoc_tool()