wrolpi/modules/docs/lib.py

import html
import logging
import pathlib
import re
from typing import List, Optional

from sqlalchemy import asc, desc, func, nullslast, text, column, bindparam, Integer, Float
from sqlalchemy.orm import Session

from wrolpi import fts
from wrolpi.collections.models import Collection
from wrolpi.collections.types import collection_type_registry
from wrolpi.errors import ValidationError
from wrolpi.files.lib import split_path_stem_and_suffix, get_mimetype, cached_search_total, \
    search_filter_cache_key

logger = logging.getLogger(__name__)

# Register author and subject collection types.
collection_type_registry.register(
    'author',
    lambda n: bool(n and n.strip()),
    'Non-empty author name'
)

collection_type_registry.register(
    'subject',
    lambda n: bool(n and n.strip()),
    'Non-empty subject name'
)

# Junk author patterns to reject.
JUNK_AUTHORS = {
    'administrator', 'admin', 'unknown', 'preferred customer', 'user',
    'owner', 'default', 'none', 'n/a', 'na', 'test', 'root',
}

JUNK_SUBJECTS = {
    'unknown', 'n/a', 'na', 'none', 'test', 'default',
}

# Regex for URL-like strings.
URL_PATTERN = re.compile(r'https?://|www\.|\.com|\.org|\.net')

# Matches year ranges like 1483-1546, -1938, 1564-1616.
YEAR_RANGE_PATTERN = re.compile(r'^-?\d{0,4}-?\d{4}$')

# Matches bracket suffixes like [from old catalog].
BRACKET_SUFFIX_PATTERN = re.compile(r'\s*\]?\s*\[.*$')

# Matches trailing birth-death years like 1843-1920 at end of string.
TRAILING_YEARS_PATTERN = re.compile(r'\s+\d{4}-\d{4}$')

# Matches pure numeric/code patterns: 00110001, 062447, 105642-01, 01/2004.
NUMERIC_CODE_PATTERN = re.compile(r'^\d[\d\-/]*$')

# Matches date strings like "14 June 2006".
DATE_STRING_PATTERN = re.compile(r'^\d{1,2}\s+\w+\s+\d{4}$')

# Matches HTML entities like • or &
HTML_ENTITY_PATTERN = re.compile(r'&\w+;')

# Matches initials pattern like J.K. (used to avoid stripping trailing period).
INITIALS_PATTERN = re.compile(r'\.\w\.')


def normalize_author(author: str) -> Optional[str]:
    """Clean an author string before validation."""
    if not author or not author.strip():
        return None
    author = author.strip()

    # Remove bracket suffixes like [from old catalog].
    author = BRACKET_SUFFIX_PATTERN.sub('', author)

    # Strip surrounding quotes.
    if len(author) >= 2 and author[0] == '"' and author[-1] == '"':
        author = author[1:-1]
    # Also strip leading quote left after bracket removal.
    if author.startswith('"'):
        author = author[1:]

    # Remove surrounding parentheses.
    if len(author) >= 2 and author[0] == '(' and author[-1] == ')':
        author = author[1:-1]

    # Remove trailing birth-death years.
    author = TRAILING_YEARS_PATTERN.sub('', author)

    # Strip trailing period, but not if it's part of initials (e.g. J.K.).
    if author.endswith('.') and not INITIALS_PATTERN.search(author):
        author = author[:-1]

    # Strip trailing comma.
    author = author.rstrip(',')

    # Title-case for consistent casing.
    author = author.title()

    author = author.strip()
    return author if author else None


def normalize_subject(subject: str) -> Optional[str]:
    """Clean a subject string before validation."""
    if not subject or not subject.strip():
        return None
    subject = subject.strip()

    # Decode HTML entities.
    subject = html.unescape(subject)

    # Title-case for consistent casing.
    subject = subject.title()

    # Truncate to 150 chars.
    if len(subject) > 150:
        subject = subject[:150]

    subject = subject.strip()
    return subject if subject else None


# Matches pure digits.
PURE_DIGITS_PATTERN = re.compile(r'^\d+$')

# Matches any letter character.
HAS_LETTER_PATTERN = re.compile(r'[a-zA-Z]')


def is_valid_author(author: str) -> bool:
    """Check if an author string looks like a real author name."""
    if not author or not author.strip():
        return False
    author = author.strip()
    if len(author) > 100:
        return False
    if len(author) < 2:
        return False
    if author.lower() in JUNK_AUTHORS:
        return False
    if URL_PATTERN.search(author):
        return False
    # Reject pure year ranges and numbers.
    if YEAR_RANGE_PATTERN.match(author):
        return False
    if PURE_DIGITS_PATTERN.match(author):
        return False
    # Reject strings with no letter characters.
    if not HAS_LETTER_PATTERN.search(author):
        return False
    return True


def is_valid_subject(subject: str) -> bool:
    """Check if a subject string looks like a real subject."""
    if not subject or not subject.strip():
        return False
    subject = subject.strip()
    if len(subject) < 2:
        return False
    if len(subject) > 150:
        return False
    if subject.lower() in JUNK_SUBJECTS:
        return False
    if URL_PATTERN.search(subject):
        return False
    # Reject pure numeric/code patterns.
    if NUMERIC_CODE_PATTERN.match(subject):
        return False
    # Reject date strings.
    if DATE_STRING_PATTERN.match(subject):
        return False
    # Reject strings with HTML entities.
    if HTML_ENTITY_PATTERN.search(subject):
        return False
    return True


# Splits multiple authors on common separators.
AUTHOR_SPLIT_PATTERN = re.compile(r'[;&]|,\s*(?:and|&)\s*|,\s+')


def split_authors(author_str: str) -> List[str]:
    """Split an author string that may contain multiple authors."""
    if not author_str:
        return []
    # Split on common separators.
    authors = AUTHOR_SPLIT_PATTERN.split(author_str)
    return [a.strip() for a in authors if a.strip()]


# Splits multiple subjects on semicolons and bare commas.
SUBJECT_SPLIT_PATTERN = re.compile(r';|,(?!\s)')


def split_subjects(subject_str: str) -> List[str]:
    """Split a subject string that may contain multiple subjects."""
    if not subject_str:
        return []
    # Split on semicolons and bare commas (comma NOT followed by a space).
    subjects = SUBJECT_SPLIT_PATTERN.split(subject_str)
    return [s.strip() for s in subjects if s.strip()]


def get_or_create_author_collection(session: Session, author_name: str) -> Optional[Collection]:
    """Find or create a Collection for an author."""
    if not author_name or not author_name.strip():
        return None

    author_name = author_name.strip()
    collection = session.query(Collection).filter(
        func.lower(Collection.name) == author_name.lower(),
        Collection.kind == 'author',
    ).first()
    if not collection:
        collection = Collection(name=author_name, kind='author')
        session.add(collection)
        session.flush()
        logger.debug(f'Created author collection: {author_name}')
    return collection


def get_or_create_subject_collection(session: Session, subject_name: str) -> Optional[Collection]:
    """Find or create a Collection for a subject."""
    if not subject_name or not subject_name.strip():
        return None

    subject_name = subject_name.strip()
    collection = session.query(Collection).filter(
        func.lower(Collection.name) == subject_name.lower(),
        Collection.kind == 'subject',
    ).first()
    if not collection:
        collection = Collection(name=subject_name, kind='subject')
        session.add(collection)
        session.flush()
        logger.debug(f'Created subject collection: {subject_name}')
    return collection


async def search_authors_by_name(session: Session, name: str, limit: int = 5) -> List[dict]:
    """Search for author collections by partial name."""
    collections = session.query(Collection) \
        .filter(Collection.kind == 'author') \
        .filter(Collection.name.ilike(f'%{name}%')) \
        .order_by(asc(Collection.name)) \
        .limit(limit) \
        .all()
    return [{'id': c.id, 'name': c.name} for c in collections]


async def search_subjects_by_name(session: Session, name: str, limit: int = 5) -> List[dict]:
    """Search for subject collections by partial name."""
    collections = session.query(Collection) \
        .filter(Collection.kind == 'subject') \
        .filter(Collection.name.ilike(f'%{name}%')) \
        .order_by(asc(Collection.name)) \
        .limit(limit) \
        .all()
    return [{'id': c.id, 'name': c.name} for c in collections]


def discover_calibre_cover(ebook_path: pathlib.Path):
    """Calibre puts a cover file in a directory which contains only one ebook (but multiple formats of the ebook),
    this function will return the path of the cover if the ebook_path is in such a situation."""
    if not ebook_path.is_file():
        raise ValueError(f'Invalid ebook path {ebook_path}')

    paths = list(ebook_path.parent.iterdir())

    ebook_path_stem, _ = split_path_stem_and_suffix(ebook_path)

    cover = None
    for path in paths:
        stem, suffix = split_path_stem_and_suffix(path)
        if path.is_dir():
            # We don't care about sub directories.
            continue
        if stem == ebook_path_stem:
            # This is one of the formats of the ebook.
            continue
        if path.name == 'metadata.opf':
            # Metadata, this is expected in a Calibre book directory.
            continue
        if stem == 'cover' and get_mimetype(path).startswith('image/'):
            # Discovered the cover.
            cover = path
            continue
        # Some file that does not match the ebook stem, metadata, or cover.  This must not be a Calibre book directory.
        return None
    return cover


def get_statistics(session: Session) -> dict:
    """Get doc statistics."""
    from .models import Doc
    from wrolpi.files.models import FileGroup

    total = session.query(func.count(Doc.id)).scalar() or 0
    total_size = session.query(func.sum(Doc.size)).scalar() or 0

    epub_count = session.query(func.count(Doc.id)).join(FileGroup).filter(
        FileGroup.mimetype.startswith('application/epub')).scalar() or 0
    pdf_count = session.query(func.count(Doc.id)).join(FileGroup).filter(
        FileGroup.mimetype == 'application/pdf').scalar() or 0

    author_count = session.query(func.count(Collection.id)).filter(
        Collection.kind == 'author').scalar() or 0
    subject_count = session.query(func.count(Collection.id)).filter(
        Collection.kind == 'subject').scalar() or 0

    return dict(statistics=dict(
        doc_count=total,
        epub_count=epub_count,
        pdf_count=pdf_count,
        other_count=total - epub_count - pdf_count,
        total_size=total_size,
        author_count=author_count,
        subject_count=subject_count,
    ))


def _doc_response(doc) -> dict:
    """Build a JSON response dict for a Doc."""
    file_group = doc.file_group.__json__()
    return {
        'file_group': file_group,
        'doc': {
            'id': doc.id,
            'publisher': doc.publisher,
            'language': doc.language,
            'page_count': doc.page_count,
            'subject': doc.subject,
            'description': doc.description,
            'size': doc.size,
        }
    }


def _get_doc(session, file_group_id: int, skip_viewed: bool = False):
    from .models import Doc
    doc = session.query(Doc).options(*Doc.json_options()).filter_by(file_group_id=file_group_id).one_or_none()
    if not doc:
        raise ValidationError(f'Doc with file_group_id {file_group_id} not found')
    if not skip_viewed:
        doc.file_group.set_viewed()
    return doc


# ORDER BY clauses for the unfiltered docs browse.  Only orderings that walk a `file_group` index
# qualify; the rest keep the ORM query.
DOC_BROWSE_ORDERS = {
    'published_datetime': 'fg.published_datetime DESC NULLS LAST, fg.id DESC',
    'id': 'fg.id DESC',
}


def _doc_browse_sql(order_by: str) -> Optional[str]:
    """SQL for a page of FileGroup ids when browsing docs with no filters, or None when `order_by` has no fast path.

    CROSS JOIN pins file_group as the outer table: SQLite walks its index in sort order and probes doc's covering
    (file_group_id) index, so no file_group row is read for the sort key."""
    order = DOC_BROWSE_ORDERS.get(order_by)
    if not order:
        return None
    return f'''SELECT fg.id AS id FROM file_group fg CROSS JOIN doc d ON d.file_group_id = fg.id
        ORDER BY {order} LIMIT :limit OFFSET :offset'''


def _search_docs(session: Session, search_str=None, author=None, subject=None, language=None, mimetype=None,
                 limit=20, offset=0, order_by='published_datetime', tag_names=None, deep=False):
    from .models import Doc
    from wrolpi.files.models import FileGroup

    # Deep search matches all FTS columns (including d_text, the document contents); the default matches only
    # the much smaller a/b/c columns.
    match = fts.translate_websearch(search_str, None if deep else fts.ABC_COLUMNS) if search_str else None

    query = session.query(FileGroup).options(*FileGroup.json_options()).join(Doc, Doc.file_group_id == FileGroup.id)

    fts_sq = None
    if match is not None:
        fts_sq = text('SELECT rowid AS id, -rank AS ts_rank'
                      ' FROM file_group_fts WHERE file_group_fts MATCH :fts_match') \
            .bindparams(fts_match=match) \
            .columns(column('id', Integer), column('ts_rank', Float)) \
            .alias('fts')
        query = query.join(fts_sq, fts_sq.c.id == FileGroup.id)

    if author:
        query = query.filter(FileGroup.b_text.ilike(f'%{author}%'))

    if subject:
        query = query.filter(Doc.subject.ilike(f'%{subject}%'))

    if language:
        query = query.filter(Doc.language == language)

    if mimetype:
        query = query.filter(FileGroup.mimetype.startswith(mimetype))

    if tag_names:
        from wrolpi.tags import Tag, TagFile
        tagged_fg_ids = session.query(TagFile.file_group_id) \
            .join(Tag, Tag.id == TagFile.tag_id) \
            .filter(Tag.name.in_(tag_names))
        query = query.filter(FileGroup.id.in_(tagged_fg_ids))

    unfiltered = not (search_str or author or subject or language or mimetype or tag_names)
    if unfiltered:
        total = cached_search_total(
            search_filter_cache_key('docs'),
            lambda: session.query(func.count(Doc.id)).scalar() or 0,
        )
    else:
        cache_key = search_filter_cache_key(
            'docs', search_str=search_str, author=author, subject=subject,
            language=language, mimetype=mimetype, tag_names=tag_names, deep=deep,
        )
        # Count before ORDER BY / LIMIT; cache by filters not page.
        total = cached_search_total(cache_key, lambda q=query: q.count())

    # Ordering.
    if order_by == 'rank' and fts_sq is not None:
        query = query.order_by(
            desc(fts_sq.c.ts_rank),
            desc(FileGroup.id),
        )
    elif order_by == 'published_datetime':
        query = query.order_by(nullslast(desc(FileGroup.published_datetime)), desc(FileGroup.id))
    elif order_by == 'size':
        query = query.order_by(nullslast(desc(Doc.size)))
    elif order_by == 'title':
        query = query.order_by(nullslast(asc(FileGroup.a_text)))
    else:
        query = query.order_by(desc(FileGroup.id))

    # `size` sorts on a doc column and `title` has no file_group index, so neither may fall back to `id`.
    browse_sql = None
    if unfiltered and order_by not in ('size', 'title'):
        browse_sql = _doc_browse_sql(order_by if order_by in DOC_BROWSE_ORDERS else 'id')
    if browse_sql:
        fg_ids = [i for i, in session.execute(text(browse_sql), dict(limit=limit, offset=offset)).fetchall()]
        by_id = {fg.id: fg for fg in session.query(FileGroup).options(*FileGroup.json_options())
                 .filter(FileGroup.id.in_(fg_ids))} if fg_ids else {}
        file_groups = [by_id[i] for i in fg_ids]
    else:
        file_groups = query.offset(offset).limit(limit).all()
    fg_ids = [fg.id for fg in file_groups]
    results = [fg.__json__() for fg in file_groups]

    if search_str and fg_ids:
        hints = _fetch_section_hints(session, fg_ids, search_str)
        for r in results:
            hint = hints.get(r['id'])
            if hint:
                r['section_hint'] = hint

    return results, total


def _fetch_section_hints(session, file_group_ids, search_str):
    """For each matching Doc, return the best-ranking DocSection for `search_str`.

    Returns a mapping of file_group_id -> {kind, ordinal, label, snippet}.
    """
    if not file_group_ids:
        return {}

    match = fts.translate_websearch(search_str)
    if match is None:
        return {}

    # FTS5 auxiliary functions (snippet, rank) must be evaluated in the innermost SELECT that scans the
    # FTS table; the window function must live in a separate outer layer.
    sql = text('''
        SELECT fg_id, kind, ordinal, label, snippet FROM (
            SELECT fg_id, kind, ordinal, label, snippet,
                   row_number() OVER (PARTITION BY fg_id
                                      ORDER BY fts_rank ASC, ordinal ASC) AS rn
            FROM (
                SELECT d.file_group_id AS fg_id,
                       ds.kind AS kind,
                       ds.ordinal AS ordinal,
                       ds.label AS label,
                       snippet(doc_section_fts, 0, '[[WROLPI_HL]]', '[[/WROLPI_HL]]', '…', 20) AS snippet,
                       doc_section_fts.rank AS fts_rank
                FROM doc_section_fts
                JOIN doc_section ds ON ds.id = doc_section_fts.rowid
                JOIN doc d ON d.id = ds.doc_id
                WHERE doc_section_fts MATCH :q AND d.file_group_id IN :ids
            )
        ) WHERE rn = 1
    ''').bindparams(bindparam('ids', expanding=True))
    rows = session.execute(sql, {'q': match, 'ids': list(file_group_ids)}).fetchall()
    hints = {}
    for row in rows:
        hints[row.fg_id] = {
            'kind': row.kind,
            'ordinal': row.ordinal,
            'label': row.label,
            'snippet': row.snippet,
        }
    return hints