wrolpi/modules/videos/lib.py

import functools
import html
import pathlib
import re
from collections import defaultdict
from typing import Tuple, Optional, Union

from sqlalchemy import or_
from sqlalchemy.orm import Session
from yt_dlp import YoutubeDL

from wrolpi import before_startup
from wrolpi.common import ConfigFile, get_media_directory, sanitize_link
from wrolpi.dates import from_timestamp, Seconds
from wrolpi.db import get_db_curs, get_db_session, optional_session
from wrolpi.files.models import File
from wrolpi.vars import PYTEST
from .common import is_valid_poster, convert_image, \
    generate_video_poster, logger, REQUIRED_OPTIONS, ConfigError, \
    get_no_channel_directory, get_video_duration
from .models import Channel, Video

logger = logger.getChild(__name__)

DEFAULT_DOWNLOAD_FREQUENCY = Seconds.week


class VideoInfoJSON(object):
    channel_source_id = None
    title = None
    duration = None
    view_count = None
    url = None


def process_video_info_json(video: Video) -> VideoInfoJSON:
    """
    Parse the Video's info json file, return the relevant data.
    """
    video_info_json = VideoInfoJSON()
    if info_json := video.get_info_json():
        title = info_json.get('fulltitle') or info_json.get('title')
        video_info_json.title = html.unescape(title) if title else None

        video_info_json.duration = info_json.get('duration')
        video_info_json.view_count = info_json.get('view_count')
        video_info_json.url = info_json.get('webpage_url') or info_json.get('url')
        video_info_json.channel_source_id = info_json.get('channel_id')

    return video_info_json


@optional_session
@functools.lru_cache
def get_channel_id_by_source_id(session: Session, source_id: str) -> Optional[int]:
    """Return the id of the Channel which matches the provided `source_id`.

    These are cached results."""
    channel = session.query(Channel).filter_by(source_id=source_id).one_or_none()
    return channel.id if channel else None


def validate_video(video: Video, channel_generate_poster: bool, session: Session):
    """
    Validate a single video.

    A Video is validated when we have attempted to find its: title, duration, view_count, url, caption, size, and
    channel.

    A Video is also valid when it has a JPEG poster, if any.  If no poster can be
    found, it will be generated from the video file.
    """
    info_json_path = video.info_json_file.path if video.info_json_file else video.info_json_path
    if info_json_path and (not video.title or not video.duration or not video.view_count or not video.url):
        # These properties can be found in the info json.
        video_info_json = process_video_info_json(video)
        video.title = video_info_json.title
        video.duration = video_info_json.duration
        video.url = video_info_json.url
        # View count will probably be overwritten by more recent data when this Video's Channel is
        # updated.
        video.view_count = video.view_count or video_info_json.view_count

        if video_info_json.channel_source_id:
            video.channel_id = get_channel_id_by_source_id(source_id=video_info_json.channel_source_id)

    video_path = video.video_path

    if not video.title or not video.upload_date or not video.source_id:
        # Video is missing things that can be extracted from the video file name.
        # These are the least trusted, so anything already on the video should be trusted.
        _, upload_date, source_id, title = parse_video_file_name(video_path)
        video.title = video.title or html.unescape(title)
        video.upload_date = video.upload_date or upload_date
        video.source_id = video.source_id or source_id

    if not video.size or not video.modification_datetime:
        stat = video_path.stat()
        video.size = video.size or stat.st_size
        video.modification_datetime = video.modification_datetime or from_timestamp(stat.st_mtime)

    if channel_generate_poster:
        # Try to convert/generate, but keep the old poster if those fail.
        new_poster_path, duration = convert_or_generate_poster(video) or video.poster_path
        if new_poster_path:
            # Poster was created/updated.
            poster_file = File.upsert(new_poster_path, session)
            session.flush([poster_file])
            video.poster_file = poster_file
            video.poster_path = new_poster_path
        if duration:
            video.duration = duration

    if not video.duration:
        # Duration was not retrieved during poster generation.
        video.duration = get_video_duration(video_path)


def convert_or_generate_poster(video: Video) -> Tuple[Optional[pathlib.Path], Optional[int]]:
    """
    If a Video has a poster, but the poster is invalid, convert it.  If a Video has no poster, generate one from the
    video file.

    Returns None if the poster was not converted, and not generated.
    """
    video_path = video.video_path
    # Modification/generation of poster is enabled for this channel.
    if video.poster_path or video.poster_file:
        # Check that the poster is a more universally supported JPEG.
        old: pathlib.Path = video.poster_path if video.poster_path else video.poster_file.path
        new = old.with_suffix('.jpg')

        if old != new and new.exists():
            # Destination JPEG already exists (it may have the wrong format).
            old.unlink()
            old = video.poster_path = new

        if not is_valid_poster(old):
            # Poster is not valid, convert it and place it in the new location.
            try:
                convert_image(old, new)
                old.unlink(missing_ok=True)
                logger.info(f'Converted invalid poster {repr(str(old))} to {repr(str(new))}')
                return new, None
            except Exception as e:
                logger.error(f'Failed to convert invalid poster {old} to {new}', exc_info=e)
                return None, None

    if not video.poster_path or not video.poster_file:
        # Video poster was not discovered, or converted.  Let's generate it.
        try:
            poster_path, duration = generate_video_poster(video_path)
            logger.debug(f'Generated poster for {video}')
            return poster_path, duration
        except Exception as e:
            logger.error(f'Failed to generate poster for {video}', exc_info=e)

    return None, None


class ChannelsConfig(ConfigFile):
    file_name = 'channels.yaml'
    default_config = dict(
        channels={
            'wrolpi': dict(
                name='WROLPi',
                url='https://www.youtube.com/channel/UC4t8bw1besFTyjW7ZBCOIrw/videos',
                directory='videos/wrolpi',
                download_frequency=604800,
            )
        },
        favorites=dict(),
    )

    @property
    def channels(self) -> dict:
        return self._config['channels']

    @channels.setter
    def channels(self, value: dict):
        self.update({'channels': value})

    @property
    def favorites(self) -> dict:
        return self._config['favorites']

    @favorites.setter
    def favorites(self, value: dict):
        self.update({'favorites': value})


CHANNELS_CONFIG: ChannelsConfig = ChannelsConfig(global_=True)
TEST_CHANNELS_CONFIG: ChannelsConfig = None


def get_channels_config():
    global TEST_CHANNELS_CONFIG
    if isinstance(TEST_CHANNELS_CONFIG, ConfigFile):
        return TEST_CHANNELS_CONFIG

    global CHANNELS_CONFIG
    return CHANNELS_CONFIG


def set_test_channels_config(enable: bool):
    global TEST_CHANNELS_CONFIG
    if enable is True:
        TEST_CHANNELS_CONFIG = ChannelsConfig()
    else:
        TEST_CHANNELS_CONFIG = None


class VideoDownloaderConfig(ConfigFile):
    file_name = 'videos_downloader.yaml'
    default_config = dict(
        continue_dl=True,
        dateafter='19900101',
        file_name_format='%(uploader)s_%(upload_date)s_%(id)s_%(title)s.%(ext)s',
        nooverwrites=True,
        quiet=False,
        writeautomaticsub=True,
        writeinfojson=True,
        writesubtitles=True,
        writethumbnail=True,
        youtube_include_dash_manifest=False,
    )

    @property
    def continue_dl(self) -> bool:
        return self._config['continue_dl']

    @continue_dl.setter
    def continue_dl(self, value: bool):
        self.update({'continue_dl': value})

    @property
    def dateafter(self) -> str:
        return self._config['dateafter']

    @dateafter.setter
    def dateafter(self, value: str):
        self.update({'dateafter': value})

    @property
    def file_name_format(self) -> str:
        return self._config['file_name_format']

    @file_name_format.setter
    def file_name_format(self, value: str):
        self.update({'file_name_format': value})

    @property
    def nooverwrites(self) -> bool:
        return self._config['nooverwrites']

    @nooverwrites.setter
    def nooverwrites(self, value: bool):
        self.update({'nooverwrites': value})

    @property
    def quiet(self) -> bool:
        return self._config['quiet']

    @quiet.setter
    def quiet(self, value: bool):
        self.update({'quiet': value})

    @property
    def writeautomaticsub(self) -> bool:
        return self._config['writeautomaticsub']

    @writeautomaticsub.setter
    def writeautomaticsub(self, value: bool):
        self.update({'writeautomaticsub': value})

    @property
    def writeinfojson(self) -> bool:
        return self._config['writeinfojson']

    @writeinfojson.setter
    def writeinfojson(self, value: bool):
        self.update({'writeinfojson': value})

    @property
    def writesubtitles(self) -> bool:
        return self._config['writesubtitles']

    @writesubtitles.setter
    def writesubtitles(self, value: bool):
        self.update({'writesubtitles': value})

    @property
    def writethumbnail(self) -> bool:
        return self._config['writethumbnail']

    @writethumbnail.setter
    def writethumbnail(self, value: bool):
        self.update({'writethumbnail': value})

    @property
    def youtube_include_dash_manifest(self) -> bool:
        return self._config['youtube_include_dash_manifest']

    @youtube_include_dash_manifest.setter
    def youtube_include_dash_manifest(self, value: bool):
        self.update({'youtube_include_dash_manifest': value})


VIDEO_DOWNLOADER_CONFIG: VideoDownloaderConfig = VideoDownloaderConfig(global_=True)
TEST_VIDEO_DOWNLOADER_CONFIG: VideoDownloaderConfig = None


def get_downloader_config() -> VideoDownloaderConfig:
    global TEST_VIDEO_DOWNLOADER_CONFIG
    if isinstance(TEST_VIDEO_DOWNLOADER_CONFIG, VideoDownloaderConfig):
        return TEST_VIDEO_DOWNLOADER_CONFIG

    global VIDEO_DOWNLOADER_CONFIG
    return VIDEO_DOWNLOADER_CONFIG


def set_test_downloader_config(enabled: bool):
    global TEST_VIDEO_DOWNLOADER_CONFIG
    if enabled:
        TEST_VIDEO_DOWNLOADER_CONFIG = VideoDownloaderConfig()
    else:
        TEST_VIDEO_DOWNLOADER_CONFIG = None


def get_channels_config_from_db(session: Session) -> dict:
    """Create a dictionary that contains all the Channels from the DB."""
    channels = session.query(Channel).order_by(Channel.directory).all()
    channels = sorted((i.config_view() for i in channels), key=lambda i: i['directory'])

    # Get all Videos that are favorites.  Store them in their own config section, so they can be preserved if a channel
    # is deleted or the DB is wiped.
    favorite_videos = session.query(Video).filter(Video.favorite != None, Video.video_path != None).all()  # noqa
    favorites = defaultdict(lambda: {})
    for video in favorite_videos:
        if video.channel:
            directory = video.channel.directory.relative_to(get_media_directory())
            favorites[str(directory)][video.video_path.name] = dict(favorite=video.favorite)
        else:
            favorites['NO CHANNEL'][video.video_path.name] = dict(favorite=video.favorite)
    favorites = dict(favorites)

    return dict(channels=channels, favorites=favorites)


def _detect_old_favorites(config: Union[ChannelsConfig, dict]) -> bool:
    # TODO remove these old favorites after beta.
    media_directory = get_media_directory()
    favorites = config.favorites if hasattr(config, 'favorites') else config['favorites']
    for channel, favorites in favorites.items():
        if channel == 'NO CHANNEL':
            continue
        if not (media_directory / channel).is_dir():
            return True
    return False


@optional_session()
def save_channels_config(session=None, preserve_favorites: bool = True):
    """Get the Channel information from the DB, save it to the config."""
    config = get_channels_config_from_db(session)
    channels_config = get_channels_config()
    # TODO remove these old favorites after beta.
    old_favorites = _detect_old_favorites(channels_config)
    if old_favorites and config['favorites']:
        # There are favorites in the database, remove the old favorites.
        channels_config.favorites = {}
    if preserve_favorites:
        config['favorites'].update(channels_config.favorites or {})
    channels_config.update(config)


@before_startup
def import_channels_config():
    """Import channel settings to the DB.  Existing channels will be updated."""
    logger.info('Importing videos config')
    try:
        config = get_channels_config()
        channels, favorites = config.channels, config.favorites

        with get_db_session(commit=True) as session:
            for data in channels:
                if isinstance(data, str):
                    # Outdated style config.
                    # TODO remove this after beta.
                    data = channels[data]
                # A Channel's directory is saved (in config) relative to the media directory.
                directory = get_media_directory() / data['directory']
                for option in (i for i in REQUIRED_OPTIONS if i not in data):
                    raise ConfigError(f'Channel "{directory}" is required to have "{option}"')

                channel = session.query(Channel).filter(
                    or_(
                        Channel.directory == str(directory),
                        Channel.directory == str(data['directory']),
                        Channel.directory == str(directory.relative_to(get_media_directory())),
                    ),
                ).one_or_none()
                if not channel:
                    # Channel not yet in the DB, add it.
                    channel = Channel(directory=directory)
                    session.add(channel)

                # Only name and directory are required
                channel.name = data['name']
                channel.directory = directory

                # A URL should not be an empty string
                data['url'] = data['url'] or None

                # Copy existing channel data, update all values from the config.  This is necessary to clear out
                # values not in the config.
                full_data = channel.dict()
                full_data.update(data)
                full_data['skip_download_videos'] = list(set(data.get('skip_download_videos', {})))
                channel.update(full_data)

                if not channel.source_id and channel.url:
                    # If we can download from a channel, we must have its source_id.
                    channel.source_id = get_channel_source_id(channel.url)

        with get_db_session(commit=True) as session:
            channels_by_link = {sanitize_link(i.name): i for i in session.query(Channel)}
            for directory_, favorites in favorites.items():
                if directory_ != 'NO CHANNEL':
                    # A Channel's directory is saved (in the config) relative to the media directory.
                    directory = get_media_directory() / directory_
                    channel: Channel = session.query(Channel).filter_by(directory=directory).one_or_none()
                    if not channel:
                        # Directory may be the outdated "link".
                        # TODO remove these old favorites after beta.
                        channel = channels_by_link.get(directory_)
                    if not channel:
                        logger.warning(f'Cannot find channel {directory=} for favorites!')
                        continue
                    channel_dir = channel.directory
                else:
                    channel_dir = get_no_channel_directory()

                # Set favorite Videos of this Channel.
                for video_path, data in favorites.items():
                    # Favorite in the config is the name of the video_path.  Add the channel directory onto this
                    # video_path, so we can match the complete path for the Video.
                    video_path = channel_dir / video_path
                    video = session.query(Video).filter_by(video_path=video_path).one_or_none()
                    # If no Video is found, it may be that we need to refresh.
                    if video:
                        video.favorite = data['favorite']
                    else:
                        logger.warning(f'Cannot find video to favorite: {video_path}')
    except Exception as e:
        logger.warning('Failed to load channels config!', exc_info=e)
        if PYTEST:
            # Do not interrupt startup, only raise during testing.
            raise


YDL = YoutubeDL()
YDL.params['logger'] = logger.getChild('youtube-dl')
YDL.add_default_info_extractors()


def get_channel_source_id(url: str) -> str:
    channel_info = YDL.extract_info(url, download=False, process=False)
    return channel_info.get('uploader_id') or channel_info['channel_id']


async def get_statistics():
    with get_db_curs() as curs:
        curs.execute('''
        SELECT
            -- total videos
            COUNT(id) AS "videos",
            -- total videos that are marked as favorite
            COUNT(id) FILTER (WHERE favorite IS NOT NULL) AS "favorites",
            -- total videos downloaded over the past week/month/year
            COUNT(id) FILTER (WHERE upload_date >= current_date - interval '1 week') AS "week",
            COUNT(id) FILTER (WHERE upload_date >= current_date - interval '1 month') AS "month",
            COUNT(id) FILTER (WHERE upload_date >= current_date - interval '1 year') AS "year",
            -- sum of all video lengths in seconds
            COALESCE(SUM(duration), 0) AS "sum_duration",
            -- sum of all video file sizes
            COALESCE(SUM(size), 0)::BIGINT AS "sum_size",
            -- largest video
            COALESCE(MAX(size), 0) AS "max_size"
        FROM
            video
        WHERE
            video_path IS NOT NULL
        ''')
        video_stats = dict(curs.fetchone())

        # Get the total videos downloaded every month for the past two years.
        curs.execute('''
        SELECT
            DATE_TRUNC('month', months.a),
            COUNT(id)::BIGINT,
            SUM(size)::BIGINT AS "size"
        FROM
            generate_series(
                date_trunc('month', current_date) - interval '2 years',
                date_trunc('month', current_date) - interval '1 month',
                '1 month'::interval) AS months(a),
            video
        WHERE
            video.upload_date >= date_trunc('month', months.a)
            AND video.upload_date < date_trunc('month', months.a) + interval '1 month'
            AND video.upload_date IS NOT NULL
            AND video.video_path IS NOT NULL
        GROUP BY
            1
        ORDER BY
            1
        ''')
        monthly_videos = [dict(i) for i in curs.fetchall()]

        historical_stats = dict(monthly_videos=monthly_videos)
        historical_stats['average_count'] = (sum(i['count'] for i in monthly_videos) // len(monthly_videos)) \
            if monthly_videos else 0
        historical_stats['average_size'] = (sum(i['size'] for i in monthly_videos) // len(monthly_videos)) \
            if monthly_videos else 0

        curs.execute('''
        SELECT
            COUNT(id) AS "channels"
        FROM
            channel
        ''')
        channel_stats = dict(curs.fetchone())
    ret = dict(statistics=dict(
        videos=video_stats,
        channels=channel_stats,
        historical=historical_stats,
    ))
    return ret


NAME_PARSER = re.compile(r'(.*?)_((?:\d+?)|(?:NA))_(?:(.{5,15})_)?(.*)\.'
                         r'(jpg|webp|flv|mp4|part|info\.json|description|webm|..\.srt|..\.vtt)')


def parse_video_file_name(video_path: pathlib.Path) -> \
        Tuple[Optional[str], Optional[str], Optional[str], Optional[str]]:
    """
    A Video's file name can have data in it, this attempts to extract what may be there.

    Example: {channel_name}_{upload_date}_{source_id}_title{ext}
    """
    video_str = str(video_path)
    if match := NAME_PARSER.match(video_str):
        channel, date, source_id, title, _ = match.groups()

        channel = None if channel == 'NA' else channel
        date = None if date == 'NA' else date

        title = title.strip()
        if date == 'NA':
            return channel, None, source_id, title
        if date is None or (len(date) == 8 and date.isdigit()):
            return channel, date, source_id, title

    # Return the stem as a last resort
    title = pathlib.Path(video_path).stem.strip()
    return None, None, None, title