#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
保留PDF的前N页
使用PyPDF2或pypdf库来处理PDF页面操作
"""

import sys
import os
from pathlib import Path

def keep_first_pages(input_pdf_path, output_pdf_path, pages_to_keep):
    """
    保留PDF的前N页

    Args:
        input_pdf_path: 输入PDF文件路径
        output_pdf_path: 输出PDF文件路径
        pages_to_keep: 要保留的页数

    Returns:
        bool: 是否成功
    """
    try:
        # 尝试使用pypdf（PyPDF2的新版本）
        try:
            from pypdf import PdfReader, PdfWriter
            reader = PdfReader(input_pdf_path)
            writer = PdfWriter()

            # 计算实际要保留的页数
            total_pages = len(reader.pages)
            actual_pages = min(pages_to_keep, total_pages)

            # 添加前N页到writer
            for i in range(actual_pages):
                writer.add_page(reader.pages[i])

            # 写入输出文件
            with open(output_pdf_path, 'wb') as output_file:
                writer.write(output_file)

            print(f"成功保留前{actual_pages}页（总共{total_pages}页）")
            return True

        except ImportError:
            # 尝试使用PyPDF2（旧版本）
            try:
                from PyPDF2 import PdfReader, PdfWriter
                reader = PdfReader(input_pdf_path)
                writer = PdfWriter()

                # 计算实际要保留的页数
                total_pages = len(reader.pages)
                actual_pages = min(pages_to_keep, total_pages)

                # 添加前N页到writer
                for i in range(actual_pages):
                    writer.add_page(reader.pages[i])

                # 写入输出文件
                with open(output_pdf_path, 'wb') as output_file:
                    writer.write(output_file)

                print(f"成功保留前{actual_pages}页（总共{total_pages}页）")
                return True

            except ImportError:
                # 尝试使用PyPDF4
                try:
                    from PyPDF4 import PdfFileReader, PdfFileWriter
                    reader = PdfFileReader(input_pdf_path)
                    writer = PdfFileWriter()

                    # 计算实际要保留的页数
                    total_pages = reader.getNumPages()
                    actual_pages = min(pages_to_keep, total_pages)

                    # 添加前N页到writer
                    for i in range(actual_pages):
                        writer.addPage(reader.getPage(i))

                    # 写入输出文件
                    with open(output_pdf_path, 'wb') as output_file:
                        writer.write(output_file)

                    print(f"成功保留前{actual_pages}页（总共{total_pages}页）")
                    return True

                except ImportError:
                    print("错误：未找到可用的PDF处理库（pypdf、PyPDF2或PyPDF4）")
                    print("请安装其中之一：pip install pypdf")
                    return False

    except Exception as e:
        print(f"处理PDF时出错：{str(e)}")
        return False

def main():
    """主函数"""
    if len(sys.argv) != 4:
        print("用法: python keep_first_pages.py <输入PDF> <输出PDF> <保留页数>")
        sys.exit(1)

    input_pdf = sys.argv[1]
    output_pdf = sys.argv[2]

    try:
        pages_to_keep = int(sys.argv[3])
        if pages_to_keep < 1:
            print("错误：保留页数必须大于0")
            sys.exit(1)
    except ValueError:
        print("错误：保留页数必须是整数")
        sys.exit(1)

    # 检查输入文件是否存在
    if not os.path.exists(input_pdf):
        print(f"错误：输入文件不存在：{input_pdf}")
        sys.exit(1)

    # 创建输出目录（如果不存在）
    output_dir = os.path.dirname(output_pdf)
    if output_dir and not os.path.exists(output_dir):
        os.makedirs(output_dir, exist_ok=True)

    # 处理PDF
    success = keep_first_pages(input_pdf, output_pdf, pages_to_keep)

    if success:
        print(f"PDF处理完成：{output_pdf}")
        sys.exit(0)
    else:
        print("PDF处理失败")
        sys.exit(1)

if __name__ == "__main__":
    main()