import re import html import ebooklib import base64 import fitz # PyMuPDF for PDF support from ebooklib import epub from bs4 import BeautifulSoup, NavigableString from PyQt5.QtWidgets import ( QDialog, QTreeWidget, QTreeWidgetItem, QDialogButtonBox, QVBoxLayout, QHBoxLayout, QTextEdit, QTreeWidgetItemIterator, QSplitter, QWidget, QPushButton, QCheckBox, QMenu, QLabel, ) from PyQt5.QtCore import Qt from abogen.utils import clean_text, calculate_text_length, detect_encoding import os import logging # Add logging import urllib.parse import markdown import textwrap # Setup logging logging.basicConfig( level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s" ) class HandlerDialog(QDialog): # Class variables to remember checkbox states between dialog instances _save_chapters_separately = False _merge_chapters_at_end = True _save_as_project = False # New class variable for save_as_project option def __init__(self, book_path, file_type=None, checked_chapters=None, parent=None): super().__init__(parent) # Determine file type if not explicitly provided if file_type: self.file_type = file_type elif book_path.lower().endswith(".pdf"): self.file_type = "pdf" elif book_path.lower().endswith((".md", ".markdown")): self.file_type = "markdown" else: self.file_type = "epub" self.book_path = book_path # Extract book name from file path book_name = os.path.splitext(os.path.basename(book_path))[0] # Set window title based on file type and book name item_type = "Chapters" if self.file_type in ["epub", "markdown"] else "Pages" self.setWindowTitle(f'Select {item_type} - {book_name}') self.resize(1200, 900) self._block_signals = False # Flag to prevent recursive signals # Configure window: remove help button and allow resizing self.setWindowFlags( Qt.Window | Qt.WindowCloseButtonHint | Qt.WindowMaximizeButtonHint ) self.setWindowModality(Qt.NonModal) # Initialize save chapters flags from class variables self.save_chapters_separately = HandlerDialog._save_chapters_separately self.merge_chapters_at_end = HandlerDialog._merge_chapters_at_end self.save_as_project = HandlerDialog._save_as_project # Load the book based on file type try: if self.file_type == "epub": self.book = epub.read_epub(book_path) elif self.file_type == "markdown": self.book = None # Markdown doesn't use ebooklib else: self.book = None except KeyError as e: logging.error( f"EPUB file is missing a referenced file: {e}. Skipping missing file." ) # Try to patch ebooklib to skip missing files (monkey-patch read_file) import types orig_read_file = None try: from ebooklib import epub as _epub_module reader_class = _epub_module.EpubReader orig_read_file = reader_class.read_file def safe_read_file(self, name): try: return orig_read_file(self, name) except KeyError: logging.warning( f"Missing file in EPUB: {name}. Returning empty bytes." ) return b"" reader_class.read_file = safe_read_file self.book = epub.read_epub(book_path) reader_class.read_file = orig_read_file # Restore except Exception as patch_e: logging.error(f"Failed to patch ebooklib for missing files: {patch_e}") raise e self.pdf_doc = fitz.open(book_path) if self.file_type == "pdf" else None self.markdown_text = None if self.file_type == "markdown": try: encoding = detect_encoding(book_path) with open(book_path, "r", encoding=encoding, errors="replace") as f: self.markdown_text = f.read() except Exception as e: logging.error(f"Error reading markdown file: {e}") self.markdown_text = "" self.markdown_toc = [] # For storing parsed markdown TOC # Extract book metadata self.book_metadata = self._extract_book_metadata() # Initialize UI elements that are used in other methods self.save_chapters_checkbox = None self.merge_chapters_checkbox = None # Build treeview self.treeWidget = QTreeWidget(self) self.treeWidget.setHeaderHidden(True) self.treeWidget.setSelectionMode(QTreeWidget.ExtendedSelection) self.treeWidget.setContextMenuPolicy(Qt.CustomContextMenu) self.treeWidget.customContextMenuRequested.connect(self.on_tree_context_menu) # Initialize checked_chapters set self.checked_chapters = set(checked_chapters) if checked_chapters else set() # For storing content and lengths self.content_texts = {} self.content_lengths = {} # Pre-process content based on file type self._preprocess_content() # Add "Information" item at the beginning of the tree info_item = QTreeWidgetItem(self.treeWidget, ["Information"]) info_item.setData(0, Qt.UserRole, "info:bookinfo") info_item.setFlags(info_item.flags() & ~Qt.ItemIsUserCheckable) font = info_item.font(0) font.setBold(True) info_item.setFont(0, font) # Build tree based on file type self._build_tree() # Hide expand/collapse decoration if there are no parent items has_parents = False for i in range(self.treeWidget.topLevelItemCount()): if self.treeWidget.topLevelItem(i).childCount() > 0: has_parents = True break self.treeWidget.setRootIsDecorated(has_parents) # Setup UI (creates save_chapters_checkbox and other UI elements) self._setup_ui() # Run auto-check after UI is setup if not self._are_provided_checks_relevant(): self._run_auto_check() # Connect signals self.treeWidget.currentItemChanged.connect(self.update_preview) self.treeWidget.itemChanged.connect(self.handle_item_check) self.treeWidget.itemChanged.connect(lambda _: self._update_checkbox_states()) self.treeWidget.itemDoubleClicked.connect(self.handle_item_double_click) # Select first item and expand all self.treeWidget.expandAll() if self.treeWidget.topLevelItemCount() > 0: self.treeWidget.setCurrentItem(self.treeWidget.topLevelItem(0)) self.treeWidget.setFocus() # Update checkbox states self._update_checkbox_states() def _preprocess_content(self): """Pre-process content from the document""" if self.file_type == "epub": try: self._process_epub_content_nav() # Use the new navigation-based method except Exception as e: logging.error( f"Error processing EPUB with navigation: {e}. Falling back to TOC/spine.", exc_info=True, ) # Fallback to a simpler spine-based processing if nav fails self._process_epub_content_spine_fallback() elif self.file_type == "markdown": self._preprocess_markdown_content() else: self._preprocess_pdf_content() def _preprocess_pdf_content(self): """Pre-process all page contents from PDF document""" for page_num in range(len(self.pdf_doc)): text = clean_text(self.pdf_doc[page_num].get_text()) # Remove bracketed numbers (citations, footnotes) text = re.sub(r"\[\s*\d+\s*\]", "", text) # Remove standalone page numbers (numbers alone on a line) text = re.sub(r"^\s*\d+\s*$", "", text, flags=re.MULTILINE) # Remove page numbers at the end of paragraphs # This pattern looks for digits surrounded by whitespace at the end of paragraphs text = re.sub(r"\s+\d+\s*$", "", text, flags=re.MULTILINE) # Also remove page numbers followed by a hyphen or dash at paragraph end # (common in headers/footers like "- 42 -") text = re.sub(r"\s+[-–—]\s*\d+\s*[-–—]?\s*$", "", text, flags=re.MULTILINE) page_id = f"page_{page_num + 1}" self.content_texts[page_id] = text self.content_lengths[page_id] = calculate_text_length(text) def _preprocess_markdown_content(self): if not self.markdown_text: return # Generate TOC from the original (dedented) markdown BEFORE cleaning, # so header ids/anchors are preserved for reliable position detection. original_text = textwrap.dedent(self.markdown_text) md = markdown.Markdown(extensions=['toc', 'fenced_code']) html = md.convert(original_text) self.markdown_toc = md.toc_tokens # Use cleaned text for stored content/length calculations cleaned_full_text = clean_text(original_text) soup = BeautifulSoup(html, 'html.parser') self.content_texts = {} self.content_lengths = {} if not self.markdown_toc: chapter_id = "markdown_content" self.content_texts[chapter_id] = cleaned_full_text self.content_lengths[chapter_id] = calculate_text_length(cleaned_full_text) return all_headers = [] def flatten_toc(toc_list): for header in toc_list: all_headers.append(header) if header.get('children'): flatten_toc(header['children']) flatten_toc(self.markdown_toc) header_positions = [] for header in all_headers: header_id = header['id'] id_pattern = f'id="{header_id}"' pos = html.find(id_pattern) if pos != -1: tag_start = html.rfind('<', 0, pos) header_positions.append({ 'id': header_id, 'start': tag_start, 'name': header['name'] }) header_positions.sort(key=lambda x: x['start']) for i, header_pos in enumerate(header_positions): header_id = header_pos['id'] header_name = header_pos['name'] content_start = header_pos['start'] content_end = header_positions[i + 1]['start'] if i + 1 < len(header_positions) else len(html) section_html = html[content_start:content_end] section_soup = BeautifulSoup(section_html, 'html.parser') header_tag = section_soup.find(attrs={'id': header_id}) if header_tag: header_tag.decompose() # Clean section text for storage/lengths section_text = clean_text(section_soup.get_text()).strip() chapter_id = header_id if section_text: full_content = f"{header_name}\n\n{section_text}" self.content_texts[chapter_id] = full_content self.content_lengths[chapter_id] = calculate_text_length(full_content) else: self.content_texts[chapter_id] = header_name self.content_lengths[chapter_id] = calculate_text_length(header_name) def _process_epub_content_spine_fallback(self): """Fallback EPUB processing based purely on spine order.""" logging.info("Using spine fallback for EPUB processing.") self.doc_content = {} spine_docs = [] for spine_item_tuple in self.book.spine: item_id = spine_item_tuple[0] item = self.book.get_item_with_id(item_id) if item: spine_docs.append(item.get_name()) else: logging.warning(f"Spine item with id '{item_id}' not found.") # Cache content for item in self.book.get_items_of_type(ebooklib.ITEM_DOCUMENT): href = item.get_name() if href in spine_docs: try: html_content = item.get_content().decode("utf-8", errors="ignore") self.doc_content[href] = html_content except Exception as e: logging.error(f"Error decoding content for {href}: {e}") self.doc_content[href] = "" # Create a simple TOC based on spine order synthetic_toc = [] self.content_texts = {} self.content_lengths = {} for i, doc_href in enumerate(spine_docs): html_content = self.doc_content.get(doc_href, "") if html_content: soup = BeautifulSoup(html_content, "html.parser") # Handle ordered lists by prepending numbers to list items for ol in soup.find_all("ol"): # Get start attribute or default to 1 start = int(ol.get("start", 1)) for i, li in enumerate(ol.find_all("li", recursive=False)): # Insert the number at the beginning of the list item number_text = f"{start + i}) " if li.string: li.string.replace_with(number_text + li.string) else: li.insert(0, NavigableString(number_text)) # Remove sup and sub tags for tag in soup.find_all(["sup", "sub"]): tag.decompose() text = clean_text(soup.get_text()).strip() if text: self.content_texts[doc_href] = text self.content_lengths[doc_href] = len(text) title = None if soup.title and soup.title.string: title = soup.title.string.strip() elif (h1 := soup.find("h1")) and h1.get_text(strip=True): title = h1.get_text(strip=True) if not title: title = f"Untitled Chapter {i + 1}" synthetic_toc.append( (epub.Link(doc_href, title, doc_href), []) ) # Wrap in tuple and empty list for compatibility # Replace book.toc with the synthetic one if it was empty or fallback was triggered if not self.book.toc or not hasattr( self, "processed_nav_structure" ): # Check if nav processing failed self.book.toc = synthetic_toc logging.info(f"Generated synthetic TOC with {len(synthetic_toc)} entries.") def _process_epub_content_nav(self): """ Process EPUB content using ITEM_NAVIGATION (NAV HTML) or ITEM_NCX. Globally orders navigation entries and slices content between them. """ logging.info( "Attempting to process EPUB using navigation document (NAV/NCX)..." ) nav_item = None nav_type = None # 1. Check ITEM_NAVIGATION for actual NAV HTML (.xhtml/.html) nav_items = list(self.book.get_items_of_type(ebooklib.ITEM_NAVIGATION)) if nav_items: # Prefer files explicitly named 'nav.xhtml' or similar preferred_nav = next( ( item for item in nav_items if "nav" in item.get_name().lower() and item.get_name().lower().endswith((".xhtml", ".html")) ), None, ) if preferred_nav: nav_item = preferred_nav nav_type = "html" logging.info(f"Found preferred NAV HTML item: {nav_item.get_name()}") else: # Check if any ITEM_NAVIGATION is actually HTML html_nav = next( ( item for item in nav_items if item.get_name().lower().endswith((".xhtml", ".html")) ), None, ) if html_nav: nav_item = html_nav nav_type = "html" logging.info( f"Found NAV HTML item in ITEM_NAVIGATION: {html_nav.get_name()}" ) # 2. If no NAV HTML found via ITEM_NAVIGATION, check if ITEM_NAVIGATION points to NCX if not nav_item and nav_items: ncx_in_nav = next( ( item for item in nav_items if item.get_name().lower().endswith(".ncx") ), None, ) if ncx_in_nav: nav_item = ncx_in_nav nav_type = "ncx" logging.info( f"Found NCX item via ITEM_NAVIGATION: {ncx_in_nav.get_name()}" ) # 3. If still no nav_item, check for NCX or fallback to NAV HTML in all ITEM_DOCUMENTs ncx_constant = getattr(epub, "ITEM_NCX", None) if not nav_item and ncx_constant is not None: ncx_items = list(self.book.get_items_of_type(ncx_constant)) if ncx_items: nav_item = ncx_items[0] nav_type = "ncx" logging.info(f"Found NCX item via ITEM_NCX: {nav_item.get_name()}") # Fallback: search all ITEM_DOCUMENTs for a NAV HTML with