import re import ebooklib import base64 import fitz # PyMuPDF for PDF support from ebooklib import epub from bs4 import BeautifulSoup, NavigableString from PyQt6.QtGui import QMovie from PyQt6.QtWidgets import ( QDialog, QTreeWidget, QTreeWidgetItem, QTextEdit, QPushButton, QVBoxLayout, QHBoxLayout, QDialogButtonBox, QSplitter, QWidget, QCheckBox, QTreeWidgetItemIterator, QLabel, QMenu, ) from PyQt6.QtCore import ( Qt, QThread, pyqtSignal, QSize, ) from abogen.utils import ( clean_text, calculate_text_length, detect_encoding, get_resource_path, ) 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 # Cache for processed book content to avoid reprocessing # Key: (book_path, modification_time, file_type) # Value: dict with content_texts, content_lengths, doc_content (for epub), markdown_toc (for markdown) _content_cache = {} class _LoaderThread(QThread): """Minimal QThread that runs a callable and emits an error string on exception.""" error = pyqtSignal(str) def __init__(self, target_callable): super().__init__() self._target = target_callable def run(self): try: self._target() except Exception as e: self.error.emit(str(e)) @classmethod def clear_content_cache(cls, book_path=None): """Clear the content cache. If book_path is provided, only clear that book's cache.""" if book_path is None: cls._content_cache.clear() logging.info("Cleared all content cache") else: keys_to_remove = [ key for key in cls._content_cache.keys() if key[0] == book_path ] for key in keys_to_remove: del cls._content_cache[key] if keys_to_remove: logging.info(f"Cleared content cache for {os.path.basename(book_path)}") def __init__(self, book_path, file_type=None, checked_chapters=None, parent=None): super().__init__(parent) # Normalize path book_path = os.path.normpath(os.path.abspath(book_path)) # 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.WindowType.Window | Qt.WindowType.WindowCloseButtonHint | Qt.WindowType.WindowMaximizeButtonHint ) self.setWindowModality(Qt.WindowModality.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.SelectionMode.ExtendedSelection) self.treeWidget.setContextMenuPolicy(Qt.ContextMenuPolicy.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 (will be filled by background loader) self.content_texts = {} self.content_lengths = {} # Add a placeholder "Information" item so the tree isn't empty immediately info_item = QTreeWidgetItem(self.treeWidget, ["Information"]) info_item.setData(0, Qt.ItemDataRole.UserRole, "info:bookinfo") info_item.setFlags(info_item.flags() & ~Qt.ItemFlag.ItemIsUserCheckable) font = info_item.font(0) font.setBold(True) info_item.setFont(0, font) # Setup UI now so dialog appears immediately self._setup_ui() # Create a centered loading overlay and show it while background load runs self._create_loading_overlay() # Hide the main UI so only the overlay is visible initially if getattr(self, "splitter", None) is not None: self.splitter.setVisible(False) self._show_loading_overlay("Loading...") # Start background loading of book content so the dialog opens immediately self._start_background_load() # 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) def _create_loading_overlay(self): """Create a centered loading indicator with a GIF on the left and text on the right. The indicator is added to the dialog's main layout above the splitter so when the splitter is hidden only the indicator is visible. """ try: # Container to hold gif + text and allow centering via stretches container = QWidget(self) container.setVisible(False) h = QHBoxLayout(container) h.setContentsMargins(0, 8, 0, 8) h.setSpacing(10) # Left: GIF label (animated) gif_label = QLabel(container) gif_label.setVisible(False) loading_gif_path = get_resource_path("abogen.assets", "loading.gif") movie = None if loading_gif_path: try: movie = QMovie(loading_gif_path) # Make GIF smaller so it doesn't dominate the text movie.setScaledSize(QSize(25, 25)) gif_label.setMovie(movie) gif_label.setFixedSize(25, 25) gif_label.setVisible(True) except Exception: movie = None # Right: Text label text_label = QLabel(container) text_label.setStyleSheet("font-size: 14pt;") # Add stretches to center the content horizontally h.addStretch(1) h.addWidget(gif_label, 0, Qt.AlignmentFlag.AlignVCenter) h.addWidget(text_label, 0, Qt.AlignmentFlag.AlignVCenter) h.addStretch(1) # Insert at top of main layout if present, otherwise keep as child try: layout = self.layout() if layout is not None: layout.insertWidget(0, container) except Exception: pass # Store refs self._loading_container = container self._loading_gif_label = gif_label self._loading_text_label = text_label self._loading_movie = movie except Exception: self._loading_container = None self._loading_gif_label = None self._loading_text_label = None self._loading_movie = None def _show_loading_overlay(self, text: str): container = getattr(self, "_loading_container", None) text_lbl = getattr(self, "_loading_text_label", None) movie = getattr(self, "_loading_movie", None) gif_lbl = getattr(self, "_loading_gif_label", None) if container is None or text_lbl is None: return text_lbl.setText(text) if movie is not None and gif_lbl is not None: try: movie.start() gif_lbl.setVisible(True) except Exception: pass container.setVisible(True) def _hide_loading_overlay(self): container = getattr(self, "_loading_container", None) movie = getattr(self, "_loading_movie", None) if container is None: return if movie is not None: try: movie.stop() except Exception: pass container.setVisible(False) def _start_background_load(self): """Start a QThread that runs the preprocessing in background.""" # Start a minimal QThread which executes _preprocess_content self._loader_thread = HandlerDialog._LoaderThread(self._preprocess_content) self._loader_thread.finished.connect(self._on_load_finished) self._loader_thread.error.connect(self._on_load_error) # ensure thread instance is deleted when done self._loader_thread.finished.connect(self._loader_thread.deleteLater) self._loader_thread.start() def _on_load_error(self, err_msg): logging.error(f"Error loading book in background: {err_msg}") if getattr(self, "previewEdit", None) is not None: self.previewEdit.setPlainText(f"Error loading book: {err_msg}") if getattr(self, "splitter", None) is not None: self.splitter.setVisible(True) self._hide_loading_overlay() def _on_load_finished(self): """Called in the main thread when background loading finished.""" # Build the tree now that content_texts/content_lengths/etc. are ready try: # Rebuild tree based on file type self._build_tree() # Run auto-check if no provided checks are relevant if not self._are_provided_checks_relevant(): self._run_auto_check() # Connect signals (after tree exists) 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) # Expand and select first item 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() # Update preview for the current selection current = self.treeWidget.currentItem() self.update_preview(current) except Exception as e: logging.error(f"Error finalizing book load: {e}") # Show the main UI and hide loading text if getattr(self, "splitter", None) is not None: self.splitter.setVisible(True) self._hide_loading_overlay() def _preprocess_content(self): """Pre-process content from the document""" # Create cache key from file path, modification time, file type, and replace_single_newlines setting try: mod_time = os.path.getmtime(self.book_path) except Exception: mod_time = 0 # Include replace_single_newlines in cache key since it affects text cleaning from abogen.utils import load_config cfg = load_config() replace_single_newlines = cfg.get("replace_single_newlines", False) cache_key = (self.book_path, mod_time, self.file_type, replace_single_newlines) # Check if content is already cached if cache_key in HandlerDialog._content_cache: cached_data = HandlerDialog._content_cache[cache_key] self.content_texts = cached_data["content_texts"] self.content_lengths = cached_data["content_lengths"] if "doc_content" in cached_data: self.doc_content = cached_data["doc_content"] if "markdown_toc" in cached_data: self.markdown_toc = cached_data["markdown_toc"] logging.info(f"Using cached content for {os.path.basename(self.book_path)}") return # Process content if not cached 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() # Cache the processed content cache_data = { "content_texts": self.content_texts, "content_lengths": self.content_lengths, } if hasattr(self, "doc_content"): cache_data["doc_content"] = self.doc_content if hasattr(self, "markdown_toc"): cache_data["markdown_toc"] = self.markdown_toc HandlerDialog._content_cache[cache_key] = cache_data logging.info(f"Cached content for {os.path.basename(self.book_path)}") 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