#!/usr/bin/env python3
"""
PDF分割脚本 - 使用pypdf，支持按书签或按单页分割。
用法: pdfsplit --bookmark 文件.pdf   或   pdfsplit --page 文件.pdf
依赖: pip install pypdf
"""

import sys
import os
import re
import subprocess
import traceback
from datetime import datetime

from pypdf import PdfReader, PdfWriter


def notify(title, message, icon="emblem-checked"):
    subprocess.run([
        "notify-send",
        "--app-name=PDF分割",
        title,
        message,
        f"--icon={icon}"
    ])


def sanitize_filename(name):
    """清理文件名中的非法字符"""
    return re.sub(r'[\\/:*?"<>|]', '_', name).strip()


def get_output_dir(pdf_path):
    """
    获取输出目录：与PDF同名的目录。
    若已存在则尝试添加(1)、(2)...后缀，直到找到不存在的目录名。
    """
    base_dir = os.path.splitext(pdf_path)[0]
    if not os.path.exists(base_dir):
        return base_dir

    n = 1
    while True:
        candidate = f"{base_dir}({n})"
        if not os.path.exists(candidate):
            return candidate
        n += 1


def _get_page_number(reader, item):
    """从书签条目中获取页码（从0开始），无法获取时返回None"""
    page = item.page if hasattr(item, 'page') else item
    if isinstance(page, int):
        return page
    try:
        return reader.get_page_number(page)
    except Exception:
        return None


def write_log(pdf_path, error):
    """在PDF同目录写入详细错误日志"""
    log_dir = os.path.dirname(pdf_path)
    basename = os.path.splitext(os.path.basename(pdf_path))[0]
    log_path = os.path.join(log_dir, f"{basename}_error.log")
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    with open(log_path, "w", encoding="utf-8") as f:
        f.write(f"PDF分割错误日志\n")
        f.write(f"时间: {timestamp}\n")
        f.write(f"文件: {pdf_path}\n")
        f.write(f"错误: {error}\n")
        f.write(f"\n详细堆栈:\n")
        f.write(traceback.format_exc())
    return log_path


def split_by_bookmark(pdf_path, output_dir):
    """按顶层书签分割PDF，每个书签覆盖的页面范围输出为独立PDF"""
    reader = PdfReader(pdf_path)
    outlines = reader.outline

    if not outlines:
        notify("未找到书签", "该PDF没有书签，无法按书签分割", "dialog-information")
        sys.exit(0)

    # 提取顶层书签：outline 可能是 OutlineItem 或 [OutlineItem, ...]
    top_bookmarks = []
    for item in outlines:
        if isinstance(item, list):
            # 嵌套列表的第一个元素通常是父书签
            top_bookmarks.append(item[0])
        else:
            top_bookmarks.append(item)

    total_pages = len(reader.pages)
    os.makedirs(output_dir, exist_ok=True)

    used_titles = {}  # 记录已用标题，处理重名
    skipped = []      # 记录被跳过的书签名称
    processed = 0

    for i, bm in enumerate(top_bookmarks):
        raw_title = bm.title if bm.title else f"书签{i + 1}"
        title = sanitize_filename(raw_title)
        if not title:
            title = f"书签{i + 1}"

        # 处理重名
        if title in used_titles:
            used_titles[title] += 1
            title = f"{title}({used_titles[title]})"
        else:
            used_titles[title] = 0

        page_num = _get_page_number(reader, bm)

        # 跳过无法获取页码或页码超出文档范围的书签
        if page_num is None or page_num >= total_pages:
            skipped.append(title)
            continue

        # 确定结束页：下一个顶层书签的起始页，或文档末尾
        if i + 1 < len(top_bookmarks):
            end_page = _get_page_number(reader, top_bookmarks[i + 1])
            # 如果下一个书签页码无效或小于等于当前页，取下一页；同时限制在文档范围内
            if end_page is None or end_page <= page_num:
                end_page = page_num + 1
            if end_page > total_pages:
                end_page = total_pages
        else:
            end_page = total_pages

        writer = PdfWriter()
        for p in range(page_num, end_page):
            writer.add_page(reader.pages[p])

        output_path = os.path.join(output_dir, f"{title}.pdf")
        writer.write(output_path)
        writer.close()
        processed += 1

    return processed, skipped


def split_by_page(pdf_path, output_dir):
    """按单页分割PDF，每页输出为独立PDF"""
    reader = PdfReader(pdf_path)
    total_pages = len(reader.pages)

    basename = os.path.splitext(os.path.basename(pdf_path))[0]
    os.makedirs(output_dir, exist_ok=True)

    for i in range(total_pages):
        writer = PdfWriter()
        writer.add_page(reader.pages[i])

        output_path = os.path.join(output_dir, f"{basename}_第{i + 1}页.pdf")
        writer.write(output_path)
        writer.close()


def main():
    try:
        if len(sys.argv) < 3:
            notify("用法: pdfsplit --bookmark|--page 文件.pdf", "", "dialog-error")
            sys.exit(1)

        if sys.argv[1] == "--bookmark":
            mode = "bookmark"
        elif sys.argv[1] == "--page":
            mode = "pages"
        else:
            notify("无效参数", "请使用 --bookmark 或 --page", "dialog-error")
            sys.exit(1)

        pdf_path = os.path.abspath(sys.argv[2])
        if not os.path.isfile(pdf_path):
            notify("文件不存在", "", "dialog-error")
            sys.exit(1)

        # 检查PDF页数，只有1页时不分割
        reader = PdfReader(pdf_path)
        if len(reader.pages) <= 1:
            notify("无需分割", "PDF文件只有1页，无需分割", "emblem-information")
            sys.exit(0)

        output_dir = get_output_dir(pdf_path)

        if mode == "bookmark":
            processed, skipped = split_by_bookmark(pdf_path, output_dir)
            if skipped:
                notify("部分书签已跳过", f"以下书签无法定位页码，已跳过: {', '.join(skipped)}", "dialog-warning")
            notify("分割成功", f"输出目录: {output_dir}\n成功分割 {processed} 个书签", "emblem-checked")
        else:
            split_by_page(pdf_path, output_dir)
            notify("分割成功", f"输出目录: {output_dir}", "emblem-checked")

    except Exception as e:
        pdf_path = os.path.abspath(sys.argv[2]) if len(sys.argv) >= 3 else ""
        log_path = write_log(pdf_path, str(e)) if pdf_path else ""
        log_name = os.path.basename(log_path) if log_path else ""
        notify("分割失败", f"详见: {log_name}" if log_name else str(e), "dialog-error")
        sys.exit(1)


if __name__ == "__main__":
    main()
