#!/usr/bin/env python

# -*- coding: utf-8 -*-

"""

PDF目录添加工具

从txt文件读取目录信息并添加到PDF文件中

"""

import re

import sys

import glob

import os

from PyPDF2 import PdfReader, PdfWriter

def parse_toc_from_txt(txt_file):

    """

    从txt文件解析目录信息

    """

    toc_items = []

    with open(txt_file, 'r', encoding='utf-8') as f:

        lines = f.readlines()

    for line in lines:

        line = line.strip()

        if not line or line == '目录正文':

            continue

        # 匹配格式：章节名称 + 点号 + 页码

        match = re.match(r'(.+?)\s*[\.。]+\s*(\d+)\s*$', line)

        if match:

            title = match.group(1).strip()

            page_num = int(match.group(2))

            toc_items.append((title, page_num))

    return toc_items

# 起始页偏移量（第一章从第8页开始，目录中写的是第1页，所以偏移量为7）

START_PAGE_OFFSET = 7

def find_pdf_file():

    """自动查找当前目录的PDF文件（排除output文件）"""

    pdf_files = glob.glob('*.pdf')

    # 排除output文件

    pdf_files = [f for f in pdf_files if 'output' not in f.lower()]

    return pdf_files[0] if pdf_files else None

def add_bookmarks_to_pdf(pdf_file, output_file, toc_items):

    """

    为PDF添加书签/目录

    """

    reader = PdfReader(pdf_file)

    writer = PdfWriter()

    # 复制所有页面

    for page in reader.pages:

        writer.add_page(page)

    # 添加书签

    for title, page_num in toc_items:

        # PDF页码从0开始，加上起始页偏移量

        pdf_page_index = page_num - 1 + START_PAGE_OFFSET

        if pdf_page_index < 0:

            pdf_page_index = 0

        elif pdf_page_index >= len(reader.pages):

            pdf_page_index = len(reader.pages) - 1

        writer.add_outline_item(title, pdf_page_index)

    # 保存输出文件

    with open(output_file, 'wb') as f:

        writer.write(f)

    print(f"成功添加 {len(toc_items)} 个目录项")

    print(f"输出文件: {output_file}")

def main():

    """主函数"""

    # 默认文件名

    txt_file = '目录.txt'

    output_file = 'output_with_toc.pdf'

    # 自动查找PDF文件

    pdf_file = find_pdf_file()

    if not pdf_file:

        print("错误: 当前目录未找到PDF文件")

        return

    # 命令行参数支持

    if len(sys.argv) >= 2:

        txt_file = sys.argv[1]

    if len(sys.argv) >= 3:

        output_file = sys.argv[2]

    print(f"读取目录文件: {txt_file}")

    print(f"读取PDF文件: {pdf_file}")

    print(f"输出文件: {output_file}")

    print("-" * 50)

    # 解析目录

    toc_items = parse_toc_from_txt(txt_file)

    if not toc_items:

        print("警告: 未找到任何目录项")

        return

    print(f"找到 {len(toc_items)} 个目录项:")

    for i, (title, page) in enumerate(toc_items, 1):

        print(f"  {i}. {title} - 第{page}页")

    print("-" * 50)

    # 添加书签到PDF

    add_bookmarks_to_pdf(pdf_file, output_file, toc_items)

if __name__ == '__main__':

    main()