#!/usr/bin/env python3
"""
PDF页面提取工具

从PDF文件中提取指定页面并保存为新PDF文件。

使用方法:
    python pdf_extractor.py input.pdf --pages "1-5,7,9-12" --output output.pdf
    python pdf_extractor.py input.pdf --pages "1-" --output output.pdf  # 提取第1页到最后
    python pdf_extractor.py input.pdf --pages "-10" --output output.pdf  # 提取前10页
"""
import argparse
import sys
from pathlib import Path
from typing import List
import io

try:
    import pypdfium2 as pdfium
    PDFIUM_AVAILABLE = True
except ImportError:
    PDFIUM_AVAILABLE = False
    pdfium = None

from loguru import logger
from .file_utils import parse_page_range


def extract_pdf_pages(
    input_pdf_path: Path,
    page_indices: List[int],
    output_pdf_path: Path
) -> bool:
    """
    从PDF中提取指定页面并保存为新PDF
    
    Args:
        input_pdf_path: 输入PDF文件路径
        page_indices: 要提取的页面索引列表（0-based）
        output_pdf_path: 输出PDF文件路径
        
    Returns:
        是否成功
    """
    if not PDFIUM_AVAILABLE:
        logger.error("❌ pypdfium2 未安装，请先安装: pip install pypdfium2")
        return False
    
    if not input_pdf_path.exists():
        logger.error(f"❌ 输入文件不存在: {input_pdf_path}")
        return False
    
    if not input_pdf_path.suffix.lower() == '.pdf':
        logger.error(f"❌ 输入文件不是PDF格式: {input_pdf_path}")
        return False
    
    try:
        # 读取PDF文件
        with open(input_pdf_path, 'rb') as f:
            pdf_bytes = f.read()
        
        # 加载PDF文档
        pdf = pdfium.PdfDocument(pdf_bytes)
        total_pages = len(pdf)
        
        if total_pages == 0:
            logger.error("❌ PDF文件为空")
            pdf.close()
            return False
        
        # 验证页面索引
        valid_indices = []
        for idx in sorted(set(page_indices)):  # 去重并排序
            if 0 <= idx < total_pages:
                valid_indices.append(idx)
            else:
                logger.warning(f"⚠️  页面索引 {idx + 1} 超出范围（总页数: {total_pages}），已跳过")
        
        if not valid_indices:
            logger.error("❌ 没有有效的页面可提取")
            pdf.close()
            return False
        
        # 创建新PDF文档
        output_pdf = pdfium.PdfDocument.new()
        
        # 导入指定页面
        success_count = 0
        for page_idx in valid_indices:
            try:
                output_pdf.import_pages(pdf, pages=[page_idx])
                success_count += 1
            except Exception as e:
                logger.warning(f"⚠️  导入第 {page_idx + 1} 页失败: {e}，已跳过")
                continue
        
        if success_count == 0:
            logger.error("❌ 没有成功导入任何页面")
            pdf.close()
            output_pdf.close()
            return False
        
        # 保存到文件
        output_pdf_path.parent.mkdir(parents=True, exist_ok=True)
        
        # 保存到内存缓冲区
        output_buffer = io.BytesIO()
        output_pdf.save(output_buffer)
        output_bytes = output_buffer.getvalue()
        
        # 写入文件
        with open(output_pdf_path, 'wb') as f:
            f.write(output_bytes)
        
        # 清理资源
        pdf.close()
        output_pdf.close()
        
        logger.info(f"✅ 成功提取 {success_count} 页到: {output_pdf_path}")
        logger.info(f"   提取的页面: {', '.join([str(idx + 1) for idx in valid_indices])}")
        return True
        
    except Exception as e:
        logger.error(f"❌ 提取PDF页面时出错: {e}")
        import traceback
        logger.debug(traceback.format_exc())
        return False


def main():
    """命令行入口"""
    parser = argparse.ArgumentParser(
        description='从PDF文件中提取指定页面并保存为新PDF文件',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例:
  # 提取第1-5页和第7页
  python pdf_extractor.py --input input.pdf --pages "1-5,7" --output output.pdf
  
  # 提取第1页到最后
  python pdf_extractor.py --input input.pdf --pages "1-" --output output.pdf
  
  # 提取前10页
  python pdf_extractor.py --input input.pdf --pages "-10" --output output.pdf
  
  # 提取单页
  python pdf_extractor.py --input input.pdf --pages "3" --output output.pdf
        """
    )
    
    parser.add_argument(
        '--input', '-i',
        type=str,
        help='输入PDF文件路径'
    )
    
    parser.add_argument(
        '--pages', '-p',
        type=str,
        required=True,
        help='要提取的页面范围，支持格式: "1-5,7,9-12", "1-", "-10", "3"'
    )
    
    parser.add_argument(
        '--output', '-o',
        type=str,
        required=True,
        help='输出PDF文件路径'
    )
    
    args = parser.parse_args()
    
    # 解析输入路径
    input_path = Path(args.input).resolve()
    
    # 解析输出路径
    output_path = Path(args.output).resolve()
    
    # 如果输出路径是目录，自动生成文件名
    if output_path.is_dir() or not output_path.suffix.lower() == '.pdf':
        output_path = output_path / f"{input_path.stem}_extracted.pdf"
    
    # 检查 pypdfium2 是否可用
    if not PDFIUM_AVAILABLE:
        logger.error("❌ pypdfium2 未安装，请先安装: pip install pypdfium2")
        sys.exit(1)
    
    # 先加载PDF获取总页数（用于验证页面范围）
    try:
        with open(input_path, 'rb') as f:
            pdf_bytes = f.read()
        pdf = pdfium.PdfDocument(pdf_bytes)
        total_pages = len(pdf)
        pdf.close()
    except Exception as e:
        logger.error(f"❌ 无法读取PDF文件: {e}")
        sys.exit(1)
    
    # 解析页面范围
    page_set = parse_page_range(args.pages, total_pages)
    page_indices = sorted(list(page_set))
    
    if not page_indices:
        logger.error(f"❌ 页面范围 '{args.pages}' 没有匹配到任何有效页面（总页数: {total_pages}）")
        sys.exit(1)
    
    logger.info(f"📋 PDF总页数: {total_pages}")
    logger.info(f"📋 要提取的页面: {args.pages} → {len(page_indices)} 页")
    
    # 执行提取
    success = extract_pdf_pages(input_path, page_indices, output_path)
    
    if success:
        logger.info(f"✅ 提取完成！输出文件: {output_path}")
        sys.exit(0)
    else:
        logger.error("❌ 提取失败")
        sys.exit(1)


if __name__ == '__main__':
    sys.exit(main())

