import os
import requests
import zipfile
import time
import json
import base64
from PIL import Image
import io
import glob
import shutil


# API 配置
API_BASE_URL = "https://jeniya.cn/v1/chat/completions"
API_KEY = "sk-5w2FNpdHmV4xoi07RAe2lWiqXGgvRWCbmCk9o8T39hZBhJIc"


def call_api(model: str, content: str) -> str:
    """
    调用API，参考NER.py的实现方式
    """
    url = API_BASE_URL
    payload = json.dumps({
        "model": model,
        "messages": [
            {
                "role": "user",
                "content": content
            }
        ]
    })
    headers = {
        'Accept': 'application/json',
        "Authorization": f"Bearer {API_KEY}",
        'Content-Type': 'application/json'
    }

    time.sleep(2)
    response = requests.request("POST", url, headers=headers, data=payload, verify=False)
    print("=============================================")
    print(f"API Response Status: {response.status_code}")
    print("=============================================")
    response_json = response.json()
    while True:
        try:
            return response_json["choices"][0]["message"]["content"]
        except:
            print(f"retry {model}")
            time.sleep(70)
            result = call_api(model, content)
            if result:
                return result


def extract_images_from_pdf_local(pdf_file_path: str, output_dir: str):
    """
    使用本地库从PDF中提取图片（基础提取）
    """
    try:
        import fitz  # PyMuPDF
    except ImportError:
        print(f"  错误: 未安装 PyMuPDF 库")
        print(f"  请运行: pip install PyMuPDF")
        return []
    
    pdf_file_name = os.path.basename(pdf_file_path)
    print(f"  使用本地库提取图片: {pdf_file_name}")
    
    jpg_files = []
    
    try:
        # 打开 PDF 文件
        pdf_document = fitz.open(pdf_file_path)
        image_count = 0
        
        # 遍历每一页
        for page_num in range(len(pdf_document)):
            page = pdf_document[page_num]
            
            # 获取页面中的图片列表
            image_list = page.get_images()
            
            for img_index, img in enumerate(image_list):
                # 获取图片的 xref
                xref = img[0]
                
                # 提取图片数据
                base_image = pdf_document.extract_image(xref)
                image_bytes = base_image["image"]
                image_ext = base_image["ext"]
                
                # 转换为 JPEG 格式
                image = Image.open(io.BytesIO(image_bytes))
                
                # 如果是 RGBA 模式，转换为 RGB
                if image.mode == 'RGBA':
                    rgb_image = Image.new('RGB', image.size, (255, 255, 255))
                    rgb_image.paste(image, mask=image.split()[3])
                    image = rgb_image
                elif image.mode != 'RGB':
                    image = image.convert('RGB')
                
                # 保存为 JPEG
                output_filename = os.path.join(output_dir, f"page_{page_num + 1}_img_{img_index + 1}.jpg")
                image.save(output_filename, "JPEG", quality=95)
                jpg_files.append(output_filename)
                image_count += 1
        
        pdf_document.close()
        print(f"  完成！共提取 {image_count} 张图片")
        return jpg_files
        
    except Exception as e:
        print(f"  本地提取错误: {e}")
        return []


def extract_images_from_pdf(pdf_file_path: str, output_dir: str, model: str = "gpt-4o-2024-11-20", use_ai_enhancement: bool = True):
    """
    从PDF文件中提取图片并保存为JPG格式
    
    Args:
        pdf_file_path: PDF文件路径
        output_dir: 输出目录（单个文件夹，如"1"）
        model: 使用的AI模型（默认使用vision模型）
        use_ai_enhancement: 是否使用AI增强图片提取（可选）
    
    Returns:
        list: 提取的JPG文件路径列表
    """
    pdf_file_name = os.path.basename(pdf_file_path)
    print(f"正在处理: {pdf_file_name}")
    
    # 步骤1: 使用本地库提取图片（直接放在output_dir中）
    print(f"步骤1: 使用本地库提取PDF中的图片...")
    jpg_files = extract_images_from_pdf_local(pdf_file_path, output_dir)
    
    if not jpg_files:
        print("  ⚠ 未提取到任何图片")
        return []
    
    print(f"  ✓ 已提取 {len(jpg_files)} 张图片")
    
    # 步骤2: 使用AI模型处理图片（可选）
    if use_ai_enhancement:
        print(f"\n步骤2: 使用AI模型 ({model}) 分析和优化图片...")
        
        prompt = f"""
你是一个专业的PDF图片提取助手。我已经从PDF文件中提取了 {len(jpg_files)} 张图片。

请分析这些图片，并确认：
1. 这些图片是否都是有效的图片（不是空白或错误）
2. 图片的质量和清晰度
3. 是否有重复或相似的图片需要去重

请以简洁的方式回复，格式如下：
- 图片总数: {len(jpg_files)}
- 有效图片: X
- 需要优化: [列出需要优化的图片编号]

如果所有图片都正常，请回复："所有图片提取成功，无需优化"
"""
        
        try:
            ai_response = call_api(model, prompt)
            print(f"  AI分析结果: {ai_response[:200]}...")
        except Exception as e:
            print(f"  ⚠ AI分析失败: {e}，继续使用本地提取的图片")
    
    print(f"\n  ✓ 完成！图片已保存到: {output_dir}")
    return jpg_files


def process_all_pdfs(model: str = "gpt-4o-2024-11-20", use_ai_enhancement: bool = True):
    """
    处理同级目录下所有PDF文件，提取图片为JPG格式
    每个PDF创建单独的一个文件夹（如"1"），所有内容都在这个文件夹中
    """
    # 获取当前工作目录
    current_dir = os.getcwd()
    print(f"正在扫描目录: {current_dir}")
    print("=" * 60)
    
    # 查找同级目录下所有PDF文件
    pdf_files = glob.glob(os.path.join(current_dir, "*.pdf"))
    
    if not pdf_files:
        print("未找到PDF文件")
        return
    
    print(f"找到 {len(pdf_files)} 个PDF文件:")
    for pdf_file in pdf_files:
        print(f"  - {os.path.basename(pdf_file)}")
    
    print("\n开始处理PDF文件...")
    print("=" * 60)
    
    # 处理每个PDF文件
    for idx, pdf_file in enumerate(pdf_files, 1):
        pdf_name_without_ext = os.path.splitext(os.path.basename(pdf_file))[0]
        
        # 为每个PDF创建一个单独的文件夹（如"1"文件夹）
        output_folder = os.path.join(current_dir, pdf_name_without_ext)
        os.makedirs(output_folder, exist_ok=True)
        print(f"\n[{idx}/{len(pdf_files)}] 处理PDF文件...")
        print(f"✓ 已创建输出文件夹: {output_folder}")
        
        # 提取图片
        jpg_files = extract_images_from_pdf(pdf_file, output_folder, model, use_ai_enhancement)
        
        if not jpg_files:
            print(f"  ⚠ 未提取到任何图片，跳过后续步骤")
            continue
        
        print(f"\n✓ PDF处理完成！共提取 {len(jpg_files)} 张JPG图片")
        
        # 步骤3: 将所有JPG文件打包为ZIP（放在同一个文件夹中）
        print("\n" + "=" * 60)
        print("步骤3: 打包所有JPG文件为ZIP...")
        
        zip_file_path = os.path.join(output_folder, "all_images.zip")
        
        print(f"正在创建ZIP文件: {os.path.basename(zip_file_path)}")
        with zipfile.ZipFile(zip_file_path, 'w', zipfile.ZIP_DEFLATED) as zip_ref:
            for jpg_file in jpg_files:
                if os.path.exists(jpg_file):  # 确保文件存在
                    # 使用文件名，不包含路径
                    arcname = os.path.basename(jpg_file)
                    zip_ref.write(jpg_file, arcname)
                    print(f"  ✓ 已添加: {arcname}")
        
        if os.path.exists(zip_file_path):
            zip_size_mb = os.path.getsize(zip_file_path) / 1024 / 1024
            print(f"✓ ZIP文件创建完成: {zip_file_path}")
            print(f"  ZIP文件大小: {zip_size_mb:.2f} MB")
        else:
            print("❌ ZIP文件创建失败")
            continue
        
        # 步骤4: 解压ZIP文件（在同一个文件夹中）
        print("\n" + "=" * 60)
        print("步骤4: 解压ZIP文件...")
        
        print(f"正在解压到: {output_folder}")
        try:
            with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
                zip_ref.extractall(output_folder)
                extracted_files = zip_ref.namelist()
            
            print(f"✓ 解压完成！共解压 {len(extracted_files)} 个文件")
        except zipfile.BadZipFile:
            print("❌ ZIP文件损坏，无法解压")
            continue
        except Exception as e:
            print(f"❌ 解压失败: {str(e)}")
            continue
        
        # 总结
        print("\n" + "=" * 60)
        print(f"处理完成总结 ({pdf_name_without_ext}):")
        print(f"  ✓ 提取JPG图片: {len(jpg_files)} 张")
        print(f"  ✓ 输出文件夹: {output_folder}")
        print(f"    - 图片文件: {output_folder}")
        print(f"    - ZIP文件: {zip_file_path}")
        print(f"    - 解压文件: {output_folder}")
        print("=" * 60)
        print()


if __name__ == "__main__":
    print("PDF图片提取工具")
    print("=" * 60)
    print("功能:")
    print("  1. 从同级目录的所有PDF文件中提取图片，保存为JPG格式")
    print("  2. 使用AI模型分析提取的图片（可选）")
    print("  3. 将所有JPG文件打包为ZIP（放在同一个文件夹中）")
    print("  4. 解压ZIP文件（在同一个文件夹中）")
    print("=" * 60)
    print()
    
    # 检查依赖
    try:
        import requests
    except ImportError:
        print("❌ 错误: 需要安装requests库")
        print("请运行: pip install requests")
        exit(1)
    
    try:
        # 可以修改这里的模型名称
        # 可用的模型参考: gpt-4o-2024-11-20, deepseek-v3, gpt-4o-mini-2024-07-18 等
        model = "gpt-4o-2024-11-20"  # 使用vision模型
        use_ai = True  # 是否使用AI增强
        
        process_all_pdfs(model=model, use_ai_enhancement=use_ai)
    except KeyboardInterrupt:
        print("\n\n⚠ 用户中断操作")
    except Exception as e:
        print(f"\n\n❌ 发生错误: {str(e)}")
        import traceback
        traceback.print_exc()