import os
import requests
import zipfile
from pathlib import Path


def extract_images_from_pdf(pdf_path, output_base_dir="processed/images"):
    """
    从PDF文件中提取图片
    
    参数:
        pdf_path: PDF文件的路径
        output_base_dir: 输出图片的基础目录，默认为 "processed/images"
    
    返回:
        tuple: (成功状态, 输出目录路径或错误信息)
    """
    try:
        # 获取PDF文件名（不含扩展名）
        pdf_filename = Path(pdf_path).stem
        
        # 创建输出目录
        output_dir = os.path.join(output_base_dir, pdf_filename)
        os.makedirs(output_dir, exist_ok=True)
        
        # 准备API请求
        api_url = "https://pdf.lutong.space/api/v1/misc/extract-images"
        
        # 打开PDF文件并准备multipart/form-data请求
        with open(pdf_path, 'rb') as pdf_file:
            files = {
                'fileInput': (os.path.basename(pdf_path), pdf_file, 'application/pdf')
            }
            data = {
                'format': 'png',
                'allowDuplicates': 'false'
            }
            headers = {
                'accept': '*/*'
            }
            
            print(f"正在处理PDF文件: {pdf_path}")
            print(f"发送请求到API...")
            
            # 发送POST请求
            response = requests.post(api_url, files=files, data=data, headers=headers)
            
            # 检查响应状态
            if response.status_code != 200:
                return False, f"API请求失败，状态码: {response.status_code}, 响应: {response.text}"
            
            print(f"成功接收到ZIP文件")
            
            # 保存ZIP文件到临时位置
            temp_zip_path = os.path.join(output_dir, f"{pdf_filename}_temp.zip")
            with open(temp_zip_path, 'wb') as zip_file:
                zip_file.write(response.content)
            
            print(f"开始解压ZIP文件...")
            
            # 解压ZIP文件
            with zipfile.ZipFile(temp_zip_path, 'r') as zip_ref:
                # 提取所有PNG文件
                for file_info in zip_ref.namelist():
                    if file_info.lower().endswith('.png'):
                        zip_ref.extract(file_info, output_dir)
                        print(f"已解压: {file_info}")
            
            # 删除临时ZIP文件
            os.remove(temp_zip_path)
            
            # 统计提取的图片数量
            png_files = [f for f in os.listdir(output_dir) if f.lower().endswith('.png')]
            print(f"成功提取 {len(png_files)} 张图片到: {output_dir}")
            
            return True, output_dir
            
    except FileNotFoundError:
        return False, f"PDF文件不存在: {pdf_path}"
    except requests.exceptions.RequestException as e:
        return False, f"网络请求错误: {str(e)}"
    except zipfile.BadZipFile:
        return False, "下载的文件不是有效的ZIP文件"
    except Exception as e:
        return False, f"处理过程中出现错误: {str(e)}"


# 使用示例
if __name__ == "__main__":
    file_path = r"F:\2.实验\3.Datasets\1.国家中小学智慧平台数据集\0.smartEdu\Course\Biology\chuyi\1\10_动物体的结构层次_视频课程\课件.pdf"
    success, result = extract_images_from_pdf(file_path)
    
    if success:
        print(f"\n✓ 成功! 图片已保存到: {result}")
    else:
        print(f"\n✗ 失败: {result}")

