wrolpi/modules/archive/models.py

import json
import pathlib
import re
from datetime import datetime
from typing import Iterable, List, Optional

import pytz
from sqlalchemy import Column, Integer, ForeignKey, BigInteger, Index
from sqlalchemy.orm import Mapped, relationship, Session, joinedload

from wrolpi import dates
from wrolpi.collections import Collection
from wrolpi.common import loaded, ModelHelper, Base, logger, get_title_from_html, get_wrolpi_config, \
    get_media_directory, strip_surrogates
from wrolpi.dates import now
from wrolpi.db import serializer
from wrolpi.errors import UnknownArchive
from wrolpi.files.models import FileGroup
from wrolpi.tags import TagFile
from wrolpi.vars import PYTEST
from .errors import InvalidArchive

logger = logger.getChild(__name__)

__all__ = ['Archive']

MATCH_URL = re.compile(r'^\s+?url:\s+?(http.*)', re.MULTILINE)
MATCH_DATE = re.compile(r'^\s+?saved date:\s+?(.*)', re.MULTILINE)


# A SingleFile header holds the page's URL, which may be very long.
SINGLEFILE_HEADER_MAX_BYTES = 64 * 1024


def read_singlefile_header(path: pathlib.Path) -> Optional[str]:
    """Return the start of a SingleFile, through the end of its header comment, or None if the header comment
    does not end within SINGLEFILE_HEADER_MAX_BYTES.  A partial header could cut off its URL or saved date."""
    head = b''
    with path.open('rb') as fh:
        # Read bytes because a compressed (SingleFileZ) singlefile has a binary ZIP tail.
        while len(head) < SINGLEFILE_HEADER_MAX_BYTES and (chunk := fh.read(1000)):
            head += chunk
            if b'-->' in head:
                return head.decode(errors='ignore')
    return None


def parse_singlefile_saved_date(saved_date: str) -> Optional[datetime]:
    """Parse SingleFile's "saved date", which is the browser's local time and its offset, into UTC.

    Returns None if the offset is malformed; the instant is unknown.

    >>> parse_singlefile_saved_date('Thu Nov 07 2024 18:36:23 GMT-0700 (Mountain Standard Time)')
    datetime.datetime(2024, 11, 8, 1, 36, 23, tzinfo=<UTC>)
    """
    parts = saved_date.split(' ')
    if len(parts) < 6:
        # No offset; assume GMT.
        return dates.strpdate(' '.join(parts[:5])).replace(tzinfo=pytz.timezone('GMT'))
    try:
        return datetime.strptime(' '.join(parts[:6]), '%a %b %d %Y %H:%M:%S GMT%z').astimezone(pytz.UTC)
    except ValueError:
        return None


class Archive(Base, ModelHelper):
    __tablename__ = 'archive'
    __table_args__ = (
        Index('idx_archive_collection_id', 'collection_id'),
    )

    id = Column(Integer, primary_key=True)

    collection_id = Column(Integer, ForeignKey('collection.id', ondelete='CASCADE'))
    collection: Mapped[Collection] = relationship('Collection', primaryjoin='Archive.collection_id==Collection.id')
    file_group_id = Column(BigInteger, ForeignKey('file_group.id', ondelete='CASCADE'), unique=True, nullable=False)
    file_group: Mapped[FileGroup] = relationship('FileGroup')

    @staticmethod
    def json_options() -> tuple:
        """Loader options for a query whose Archives will be serialized: `__json__` reads the FileGroup and its Tags,
        and the domain (the Collection's name)."""
        return joinedload(Archive.file_group).selectinload(FileGroup.tag_files), joinedload(Archive.collection)

    @serializer
    def __json__(self) -> dict:
        d = self.file_group.__json__()
        d['archive'] = dict(
            id=self.id,
            collection_id=self.collection_id,
            domain=self.domain,
            singlefile_path=self.singlefile_path.name if self.singlefile_path else None,
            readability_txt_path=self.readability_txt_path.name if self.readability_txt_path else None,
            screenshot_path=self.screenshot_path.name if self.screenshot_path else None,
        )
        return d

    def __repr__(self):
        id_, url = loaded(self, 'id'), loaded(self, 'file_group', 'url')
        path = repr(str(loaded(self, 'file_group', 'primary_path')))
        if domain_name := loaded(self, 'collection', 'name'):
            return f'<Archive id={id_} url={url} path={path} domain={domain_name}>'
        return f'<Archive id={id_} url={url} path={path}>'

    @property
    def domain(self) -> str | None:
        """
        Get the domain name for this Archive.

        This property provides backward compatibility by returning the Collection's name
        (which is the domain string for domain collections).

        Returns:
            The domain string (e.g., "example.com") or None if no collection
        """
        if self.collection and self.collection.kind == 'domain':
            return self.collection.name
        return None

    @staticmethod
    def get_by_id(session: Session, id_: int) -> Optional['Archive']:
        """Attempt to find an Archive with the provided id.  Returns None if it cannot be found."""
        archive = session.query(Archive).filter_by(id=id_).one_or_none()
        return archive

    @staticmethod
    def find_by_id(session: Session, id_: int) -> 'Archive':
        """Find an Archive with the provided id, raises an exception when no Archive is found.

        @raise UnknownArchive: if the Archive can not be found"""
        archive = Archive.get_by_id(session, id_)
        if not archive:
            raise UnknownArchive(f'Cannot find Archive with id {id_}')
        return archive

    @staticmethod
    def get_by_path(session: Session, path: pathlib.Path | str) -> Optional['Archive']:
        archive = session.query(Archive).join(FileGroup).filter(FileGroup.primary_path == str(path)).one_or_none()
        return archive

    @staticmethod
    def find_by_path(session: Session, path: pathlib.Path | str) -> Optional['Archive']:
        archive = Archive.get_by_path(session, path)
        if archive:
            return archive
        raise UnknownArchive(f'Cannot find Archive with path: {path}')

    @staticmethod
    def can_model(file_group: FileGroup) -> bool:
        # Local import to avoid circular import within archive module
        from modules.archive.lib import is_singlefile_file
        if file_group.mimetype.startswith('text') and is_singlefile_file(file_group.primary_path):
            return True
        return False

    @staticmethod
    def do_model(session: Session, file_group: FileGroup) -> 'Archive':
        # Local import to avoid circular import within archive module
        from modules.archive import model_archive
        archive = model_archive(session, file_group)
        archive.validate()
        file_group.indexed = True
        return archive

    def my_paths(self, *mimetypes: str) -> List[pathlib.Path]:
        return self.file_group.my_paths(*mimetypes)

    def my_files(self, *mimetypes: str) -> List[dict]:
        return self.file_group.my_files(*mimetypes)

    def unlink(self):
        """
        Remove any files in this archive.
        """
        # Don't fail if a file is missing, it could have been deleted manually.
        for path in self.my_paths():
            path.unlink(missing_ok=True)

    def __gt__(self, other) -> bool:
        """
        Compare Archives according to their archive datetime.
        """
        if not isinstance(other, Archive):
            raise ValueError(f'Cannot compare {type(other)} to Archive!')

        return self.file_group.published_datetime > other.file_group.published_datetime

    def delete(self, force: bool = False):
        self.file_group.delete(force=force)

        session = Session.object_session(self)
        collection = self.collection
        session.delete(self)

        if collection and collection.kind == 'domain':
            # Delete a domain collection if it has no Archives.
            remaining_archives = session.query(Archive).filter_by(collection_id=collection.id).count()
            if remaining_archives == 0:
                session.delete(collection)

    @property
    def history(self) -> Iterable[Base]:
        """Get a list of Archives that share my URL."""
        session = Session.object_session(self)
        history = list(session.query(Archive).filter(
            Archive.id != self.id,
            FileGroup.url == self.file_group.url,
            FileGroup.url != None,  # noqa
        ) \
                       .join(FileGroup, FileGroup.id == Archive.file_group_id) \
                       .options(*Archive.json_options()) \
                       .order_by(FileGroup.published_datetime))
        return history

    @property
    def singlefile_file(self) -> Optional[dict]:
        # Local import to avoid circular import within archive module
        from modules.archive.lib import is_singlefile_file
        for file in self.file_group.my_html_files():
            if is_singlefile_file(file['path']):
                return file

    @property
    def singlefile_path(self) -> Optional[pathlib.Path]:
        if singlefile_file := self.singlefile_file:
            return singlefile_file['path']

    @property
    def readability_file(self) -> Optional[dict]:
        files = self.file_group.my_files('text/html')
        for file in files:
            if file['path'].name.endswith('.readability.html'):
                return file

    @property
    def readability_path(self) -> Optional[pathlib.Path]:
        if readability_file := self.readability_file:
            return readability_file['path']

    @property
    def readability_txt_file(self) -> Optional[dict]:
        files = self.file_group.my_files('text/plain', 'text/html')
        for file in files:
            if file['path'].name.endswith('.readability.txt'):
                return file

    @property
    def readability_txt_path(self) -> Optional[pathlib.Path]:
        if readability_txt_file := self.readability_txt_file:
            return readability_txt_file['path']

    @property
    def readability_json_file(self) -> Optional[dict]:
        files = self.file_group.my_files()
        for file in files:
            if file['path'].name.endswith('.readability.json'):
                return file

    @property
    def readability_json_path(self) -> Optional[pathlib.Path]:
        if readability_json_file := self.readability_json_file:
            return readability_json_file['path']

    @property
    def screenshot_file(self) -> dict | None:
        files = self.file_group.my_files('image/')
        for file in files:
            return file

    @property
    def screenshot_path(self) -> Optional[pathlib.Path]:
        if screenshot_file := self.screenshot_file:
            return screenshot_file['path']

    def set_screenshot(self, screenshot_path: pathlib.Path) -> None:
        """Set the screenshot for this Archive.

        Args:
            screenshot_path: Path to the screenshot file

        Raises:
            ValueError: If Archive already has a screenshot or path doesn't exist
        """
        # Check if screenshot already exists
        if self.screenshot_path:
            raise ValueError(f'Archive {self.id} already has a screenshot at {self.screenshot_path}')

        # Validate path exists
        if not screenshot_path.is_file():
            raise ValueError(f'Screenshot path does not exist: {screenshot_path}')

        # Add to FileGroup.files
        self.file_group.append_files(screenshot_path)

        # Update FileGroup.data for quick lookup (similar to ebook cover_path pattern)
        # Store as filename only (relative to directory), resolved via file_group.resolve_path()
        data = dict(self.file_group.data) if self.file_group.data else {}
        data['screenshot_path'] = screenshot_path.name  # Store filename only
        self.file_group.data = data

        # Validate the archive
        self.validate()

    @property
    def info_json_file(self) -> Optional[dict]:
        files = self.file_group.my_files()
        for file in files:
            if file['path'].name.endswith('.info.json'):
                return file

    @property
    def info_json_path(self) -> Optional[pathlib.Path]:
        if info_json_file := self.info_json_file:
            return info_json_file['path']

    def apply_readability_data(self):
        """Read the Readability JSON file, apply its contents to this record."""
        readability_json_path = self.readability_json_path
        if not readability_json_path:
            logger.debug(f'{self.singlefile_path} does not have an info json file')
            return
        if not readability_json_path.is_file():
            error = f'Cannot read data from {readability_json_path} because it not exist.'
            if PYTEST:
                raise ValueError(error)
            logger.error(error)
            return

        try:
            with readability_json_path.open() as fh:
                json_contents = json.load(fh)
                # Scraped strings can contain unpaired surrogates which sqlite3 refuses to store.
                url = strip_surrogates(json_contents.get('url'))
                title = strip_surrogates(json_contents.get('title'))
        except Exception as e:
            raise InvalidArchive() from e

        # Readability is most trusted, it should overwrite any previous data.
        self.file_group.url = url or self.file_group.url
        self.file_group.title = title or self.file_group.title

    def apply_singlefile_data(self):
        """Read the start of the singlefile (if any) and extract any Archive information."""
        path = self.singlefile_path
        if not path:
            # Can't read contents of nothing.
            raise ValueError('Cannot read singlefile data when this has no files')

        if self.file_group.url and self.file_group.published_datetime:
            # This data has already been read.
            return

        head = read_singlefile_header(path)
        if not head or 'Page saved with SingleFile' not in head:
            logger.error(f'Could not find SingleFile header in {self}')
            return

        if not self.file_group.url:
            try:
                if match := MATCH_URL.findall(head):
                    self.file_group.url = match[0].strip()
            except Exception as e:
                logger.error(f'Could not get URL from singlefile {path}', exc_info=e)

        if not self.file_group.download_datetime:
            try:
                if match := MATCH_DATE.findall(head):
                    self.file_group.download_datetime = parse_singlefile_saved_date(match[0].strip())
            except Exception as e:
                logger.error(f'Could not get archive date from singlefile {path}', exc_info=e)

    def apply_domain(self):
        """Get the domain collection from the URL."""
        # Local import to avoid circular import within archive module
        from modules.archive.lib import get_or_create_domain_collection
        collection = None
        if self.file_group.url:
            session = Session.object_session(self)
            if not session:
                raise ValueError('No session found!')
            collection = get_or_create_domain_collection(session, self.file_group.url)
        # Clear collection if the URL is missing.
        self.collection_id = collection.id if collection else None

    def apply_singlefile_title(self):
        """Get the title from the Singlefile, if it's missing."""
        # Local import to avoid circular import within archive module
        from modules.archive import lib
        if self.singlefile_path and not self.file_group.title:
            self.file_group.title = strip_surrogates(get_title_from_html(lib.read_singlefile_html(self.singlefile_path)))

    def apply_metadata(self):
        """Read and apply <meta> (and more) data from the Singlefile HTML."""
        # Local import to avoid circular import within archive module
        from modules.archive import lib
        contents = lib.read_singlefile_html(self.singlefile_path)

        metadata = lib.parse_article_html_metadata(contents)
        if metadata.author:
            self.file_group.author = strip_surrogates(metadata.author)
        if metadata.title:
            self.file_group.title = strip_surrogates(metadata.title)
        if metadata.description:
            self.file_group.b_text = strip_surrogates(metadata.description)
        if metadata.published_datetime:
            self.file_group.published_datetime = metadata.published_datetime
        if metadata.modified_datetime:
            self.file_group.published_modified_datetime = metadata.modified_datetime

    def rebuild_data(self):
        """Rebuild the data field with current file paths.

        Called during validate() to ensure data field reflects actual filenames
        after reorganization or other file operations.
        """
        self.file_group.data = {
            'id': self.id,
            'domain': self.domain,
            'readability_json_path': self.readability_json_path.name if self.readability_json_path else None,
            'readability_path': self.readability_path.name if self.readability_path else None,
            'readability_txt_path': self.readability_txt_path.name if self.readability_txt_path else None,
            'screenshot_path': self.screenshot_path.name if self.screenshot_path else None,
            'singlefile_path': self.singlefile_path.name if self.singlefile_path else None,
            'info_json_path': self.info_json_path.name if self.info_json_path else None,
        }

    def validate(self):
        """Fill in any missing data about this Archive from its files."""
        try:
            self.apply_readability_data()
            self.apply_singlefile_data()
            self.apply_domain()
            self.apply_singlefile_title()
            self.apply_metadata()
            self.rebuild_data()
        except Exception as e:
            logger.warning(f'Unable to validate {self}', exc_info=e)
            if PYTEST:
                raise

    @staticmethod
    def from_paths(session: Session, *paths: pathlib.Path) -> 'Archive':
        """Create a new Archive and FileGroup from the provided paths.

        The files will be read and Archive data extracted."""
        # Local import to avoid circular import within archive module
        from modules.archive import model_archive
        file_group = FileGroup.from_paths(session, *paths)
        archive = model_archive(session, file_group)
        return archive

    def add_tag(self, session: Session, tag_id_or_name: int | str) -> TagFile:
        return self.file_group.add_tag(session, tag_id_or_name)

    @property
    def location(self):
        """The location where this Archive can be viewed in the UI."""
        return f'/archives/{self.file_group_id}'

    @property
    def download_directory(self) -> pathlib.Path:
        """
        Get the download directory for this Archive based on its domain and date.

        This uses the configured archive_destination template with variables:
        - domain: the domain name from the Collection
        - domain_tag: the tag name from the Collection (empty string if no tag)
        - year, month, day: current date

        Returns:
            Absolute path where new archives for this domain should be downloaded
        """
        archive_destination = get_wrolpi_config().archive_destination

        # Get the tag name from the collection if it exists
        domain_tag = ''
        if self.collection and self.collection.tag:
            domain_tag = self.collection.tag.name

        now_ = now()
        variables = dict(
            domain=self.domain or 'unknown',
            domain_tag=domain_tag,
            year=now_.year,
            month=now_.month,
            day=now_.day,
        )
        archive_destination = archive_destination % variables
        archive_destination = pathlib.Path(archive_destination.lstrip('/'))
        if not archive_destination.is_absolute():
            archive_destination = get_media_directory() / archive_destination

        return archive_destination