reader3.py

"""
Parses an EPUB file into a structured object that can be used to serve the book via a web interface.
"""

import os
import pickle
import shutil
from dataclasses import dataclass, field
from typing import List, Dict, Optional, Any
from datetime import datetime
from urllib.parse import unquote

import ebooklib
from ebooklib import epub
from bs4 import BeautifulSoup, Comment

# --- Data structures ---

@dataclass
class ChapterContent:
    """
    Represents a physical file in the EPUB (Spine Item).
    A single file might contain multiple logical chapters (TOC entries).
    """
    id: str           # Internal ID (e.g., 'item_1')
    href: str         # Filename (e.g., 'part01.html')
    title: str        # Best guess title from file
    content: str      # Cleaned HTML with rewritten image paths
    text: str         # Plain text for search/LLM context
    order: int        # Linear reading order


@dataclass
class TOCEntry:
    """Represents a logical entry in the navigation sidebar."""
    title: str
    href: str         # original href (e.g., 'part01.html#chapter1')
    file_href: str    # just the filename (e.g., 'part01.html')
    anchor: str       # just the anchor (e.g., 'chapter1'), empty if none
    children: List['TOCEntry'] = field(default_factory=list)


@dataclass
class BookMetadata:
    """Metadata"""
    title: str
    language: str
    authors: List[str] = field(default_factory=list)
    description: Optional[str] = None
    publisher: Optional[str] = None
    date: Optional[str] = None
    identifiers: List[str] = field(default_factory=list)
    subjects: List[str] = field(default_factory=list)


@dataclass
class Book:
    """The Master Object to be pickled."""
    metadata: BookMetadata
    spine: List[ChapterContent]  # The actual content (linear files)
    toc: List[TOCEntry]          # The navigation tree
    images: Dict[str, str]       # Map: original_path -> local_path

    # Meta info
    source_file: str
    processed_at: str
    cover_image: Optional[str] = None  # Cover image filename
    version: str = "3.0"


# --- Utilities ---

def clean_html_content(soup: BeautifulSoup) -> BeautifulSoup:

    # Remove dangerous/useless tags
    for tag in soup(['script', 'style', 'iframe', 'video', 'nav', 'form', 'button']):
        tag.decompose()

    # Remove HTML comments
    for comment in soup.find_all(string=lambda text: isinstance(text, Comment)):
        comment.extract()

    # Remove input tags
    for tag in soup.find_all('input'):
        tag.decompose()

    return soup


def extract_plain_text(soup: BeautifulSoup) -> str:
    """Extract clean text for LLM/Search usage."""
    text = soup.get_text(separator=' ')
    # Collapse whitespace
    return ' '.join(text.split())


def parse_toc_recursive(toc_list, depth=0) -> List[TOCEntry]:
    """
    Recursively parses the TOC structure from ebooklib.
    """
    result = []

    for item in toc_list:
        # ebooklib TOC items are either `Link` objects or tuples (Section, [Children])
        if isinstance(item, tuple):
            section, children = item
            entry = TOCEntry(
                title=section.title,
                href=section.href,
                file_href=section.href.split('#')[0],
                anchor=section.href.split('#')[1] if '#' in section.href else "",
                children=parse_toc_recursive(children, depth + 1)
            )
            result.append(entry)
        elif isinstance(item, epub.Link):
            entry = TOCEntry(
                title=item.title,
                href=item.href,
                file_href=item.href.split('#')[0],
                anchor=item.href.split('#')[1] if '#' in item.href else ""
            )
            result.append(entry)
        # Note: ebooklib sometimes returns direct Section objects without children
        elif isinstance(item, epub.Section):
             entry = TOCEntry(
                title=item.title,
                href=item.href,
                file_href=item.href.split('#')[0],
                anchor=item.href.split('#')[1] if '#' in item.href else ""
            )
             result.append(entry)

    return result


def get_fallback_toc(book_obj) -> List[TOCEntry]:
    """
    If TOC is missing, build a flat one from the Spine.
    """
    toc = []
    for item in book_obj.get_items():
        if item.get_type() == ebooklib.ITEM_DOCUMENT:
            name = item.get_name()
            # Try to guess a title from the content or ID
            title = item.get_name().replace('.html', '').replace('.xhtml', '').replace('_', ' ').title()
            toc.append(TOCEntry(title=title, href=name, file_href=name, anchor=""))
    return toc


def extract_metadata_robust(book_obj) -> BookMetadata:
    """
    Extracts metadata handling both single and list values.
    """
    def get_list(key):
        data = book_obj.get_metadata('DC', key)
        return [x[0] for x in data] if data else []

    def get_one(key):
        data = book_obj.get_metadata('DC', key)
        return data[0][0] if data else None

    return BookMetadata(
        title=get_one('title') or "Untitled",
        language=get_one('language') or "en",
        authors=get_list('creator'),
        description=get_one('description'),
        publisher=get_one('publisher'),
        date=get_one('date'),
        identifiers=get_list('identifier'),
        subjects=get_list('subject')
    )


# --- Main Conversion Logic ---

def process_epub(epub_path: str, output_dir: str) -> Book:

    # 1. Load Book
    print(f"Loading {epub_path}...")
    book = epub.read_epub(epub_path)

    # 2. Extract Metadata
    metadata = extract_metadata_robust(book)

    # 3. Prepare Output Directories
    if os.path.exists(output_dir):
        shutil.rmtree(output_dir)
    images_dir = os.path.join(output_dir, 'images')
    os.makedirs(images_dir, exist_ok=True)

    # 4. Extract Images & Build Map (including cover)
    print("Extracting images...")
    image_map = {} # Key: internal_path, Value: local_relative_path
    cover_image = None

    # Try to find cover image from metadata
    cover_item = None
    
    # Method 1: Check for ITEM_COVER type (most reliable)
    for item in book.get_items():
        if item.get_type() == ebooklib.ITEM_COVER:
            cover_item = item
            print(f"✓ Found cover (type COVER): {item.get_name()}")
            break
    
    # Method 2: Look for images with 'cover' or 'cvi' in the name
    if not cover_item:
        for item in book.get_items():
            if item.get_type() in (ebooklib.ITEM_IMAGE, ebooklib.ITEM_COVER):
                name_lower = item.get_name().lower()
                if 'cover' in name_lower or 'cvi' in name_lower:
                    cover_item = item
                    print(f"✓ Found cover (by name): {item.get_name()}")
                    break
    
    # Method 3: Use first large image as fallback (skip small icons/logos)
    if not cover_item:
        for item in book.get_items():
            if item.get_type() in (ebooklib.ITEM_IMAGE, ebooklib.ITEM_COVER):
                # Skip very small images (likely icons)
                if len(item.get_content()) > 10000:  # > 10KB
                    cover_item = item
                    print(f"✓ Using first large image as cover: {item.get_name()}")
                    break

    saved_files = {}  # Track saved filenames to detect collisions
    
    for item in book.get_items():
        # Extract both ITEM_IMAGE and ITEM_COVER types
        if item.get_type() in (ebooklib.ITEM_IMAGE, ebooklib.ITEM_COVER):
            # Normalize filename
            original_fname = os.path.basename(item.get_name())
            # Sanitize filename for OS
            safe_fname = "".join([c for c in original_fname if c.isalpha() or c.isdigit() or c in '._-']).strip()

            # Handle filename collisions by adding a counter
            if safe_fname in saved_files:
                base, ext = os.path.splitext(safe_fname)
                counter = 1
                while f"{base}_{counter}{ext}" in saved_files:
                    counter += 1
                safe_fname = f"{base}_{counter}{ext}"
                print(f"Warning: Filename collision, renamed to {safe_fname}")

            # Save to disk
            local_path = os.path.join(images_dir, safe_fname)
            with open(local_path, 'wb') as f:
                f.write(item.get_content())
            
            saved_files[safe_fname] = item.get_name()

            # Map keys: We try both the full internal path and just the basename
            # to be robust against messy HTML src attributes
            rel_path = f"images/{safe_fname}"
            image_map[item.get_name()] = rel_path
            image_map[original_fname] = rel_path
            
            # Check if this is the cover image
            if cover_item and item.get_name() == cover_item.get_name():
                cover_image = safe_fname

    # 5. Process TOC
    print("Parsing Table of Contents...")
    toc_structure = parse_toc_recursive(book.toc)
    if not toc_structure:
        print("Warning: Empty TOC, building fallback from Spine...")
        toc_structure = get_fallback_toc(book)

    # 6. Process Content (Spine-based to preserve HTML validity)
    print("Processing chapters...")
    spine_chapters = []

    # We iterate over the spine (linear reading order)
    for i, spine_item in enumerate(book.spine):
        item_id, linear = spine_item
        item = book.get_item_with_id(item_id)

        if not item:
            continue

        if item.get_type() == ebooklib.ITEM_DOCUMENT:
            # Raw content
            raw_content = item.get_content().decode('utf-8', errors='ignore')
            soup = BeautifulSoup(raw_content, 'html.parser')

            # A. Fix Images
            for img in soup.find_all('img'):
                src = img.get('src', '')
                if not src: continue

                # Decode URL (part01/image%201.jpg -> part01/image 1.jpg)
                src_decoded = unquote(src)
                filename = os.path.basename(src_decoded)

                # Try to find in map
                if src_decoded in image_map:
                    img['src'] = image_map[src_decoded]
                elif filename in image_map:
                    img['src'] = image_map[filename]

            # B. Clean HTML
            soup = clean_html_content(soup)

            # C. Extract Body Content only
            body = soup.find('body')
            if body:
                # Extract inner HTML of body
                final_html = "".join([str(x) for x in body.contents])
            else:
                final_html = str(soup)

            # D. Create Object
            chapter = ChapterContent(
                id=item_id,
                href=item.get_name(), # Important: This links TOC to Content
                title=f"Section {i+1}", # Fallback, real titles come from TOC
                content=final_html,
                text=extract_plain_text(soup),
                order=i
            )
            spine_chapters.append(chapter)

    # 7. Final Assembly
    final_book = Book(
        metadata=metadata,
        spine=spine_chapters,
        toc=toc_structure,
        images=image_map,
        source_file=os.path.basename(epub_path),
        processed_at=datetime.now().isoformat(),
        cover_image=cover_image
    )

    return final_book


def save_to_pickle(book: Book, output_dir: str):
    p_path = os.path.join(output_dir, 'book.pkl')
    with open(p_path, 'wb') as f:
        pickle.dump(book, f)
    print(f"Saved structured data to {p_path}")


# --- CLI ---

def sanitize_folder_name(name: str) -> str:
    """
    Sanitize folder name while preserving Unicode characters (including Chinese).
    Only removes characters that are invalid for Windows/Unix filesystems.
    """
    # Characters not allowed in Windows filenames
    invalid_chars = '<>:"/\\|?*'
    for char in invalid_chars:
        name = name.replace(char, '_')
    
    # Remove leading/trailing spaces and dots
    name = name.strip('. ')
    
    # Limit length to avoid path issues (Windows has 260 char limit)
    if len(name) > 100:
        name = name[:100]
    
    return name


if __name__ == "__main__":

    import sys
    if len(sys.argv) < 2:
        print("Usage: python reader3.py <file.epub>")
        sys.exit(1)

    epub_file = sys.argv[1]
    assert os.path.exists(epub_file), "File not found."
    
    # Create books directory if it doesn't exist
    books_dir = "books"
    os.makedirs(books_dir, exist_ok=True)
    
    # First, do a quick metadata extraction to get the real title
    print(f"Reading metadata from {epub_file}...")
    temp_book = epub.read_epub(epub_file)
    temp_metadata = extract_metadata_robust(temp_book)
    
    # Use the actual book title for folder name (supports Chinese!)
    book_title = temp_metadata.title or os.path.splitext(os.path.basename(epub_file))[0]
    safe_title = sanitize_folder_name(book_title)
    out_dir = os.path.join(books_dir, safe_title)
    
    # If folder exists, add a number suffix
    if os.path.exists(out_dir):
        counter = 1
        while os.path.exists(f"{out_dir}_{counter}"):
            counter += 1
        out_dir = f"{out_dir}_{counter}"
    
    print(f"Output directory: {out_dir}")

    book_obj = process_epub(epub_file, out_dir)
    save_to_pickle(book_obj, out_dir)
    
    # Use safe printing to avoid Unicode errors on Windows
    try:
        print("\n--- Summary ---")
        print(f"Title: {book_obj.metadata.title}")
        print(f"Authors: {', '.join(book_obj.metadata.authors)}")
        print(f"Physical Files (Spine): {len(book_obj.spine)}")
        print(f"TOC Root Items: {len(book_obj.toc)}")
        print(f"Images extracted: {len(book_obj.images)}")
        print(f"\nBook data saved to: {out_dir}")
    except UnicodeEncodeError:
        # Fallback for Windows console encoding issues
        print("\n--- Summary ---")
        print(f"Title: [Unicode title]")
        print(f"Authors: [Unicode authors]")
        print(f"Physical Files (Spine): {len(book_obj.spine)}")
        print(f"TOC Root Items: {len(book_obj.toc)}")
        print(f"Images extracted: {len(book_obj.images)}")
        print(f"\nBook data saved to: {out_dir}")