wrolpi/modules/videos/common.py

import json
import os
import pathlib
import re
import subprocess
import tempfile
from datetime import timedelta
from decimal import Decimal
from functools import partial
from pathlib import Path
from typing import List, Optional, Tuple

import PIL
from PIL import Image
from sqlalchemy.orm import Session

from wrolpi.captions import FFMPEG_BIN
from wrolpi.cmd import FFPROBE_BIN
from wrolpi.common import logger, iterify, get_media_directory, \
    minimize_dict, match_paths_to_suffixes
from wrolpi.db import get_db_session, get_db_curs
from wrolpi.errors import ChannelNameConflict, ChannelURLConflict, \
    ChannelDirectoryConflict, ChannelSourceIdConflict
from wrolpi.files.models import File
from wrolpi.vars import DEFAULT_FILE_PERMISSIONS
from .models import Channel

logger = logger.getChild(__name__)

REQUIRED_OPTIONS = ['name', 'directory']
MINIMUM_CHANNEL_KEYS = {'id', 'name', 'directory', 'url', 'video_count'}
MINIMUM_INFO_JSON_KEYS = {'description', 'view_count', 'webpage_url'}
MINIMUM_VIDEO_KEYS = {'id', 'title', 'upload_date', 'duration', 'channel', 'channel_id', 'favorite', 'size',
                      'poster_path', 'caption_path', 'video_path', 'info_json', 'channel', 'viewed', 'source_id',
                      'view_count'}
# These are the supported video formats.  These are in order of their preference.
VIDEO_EXTENSIONS = ('mp4', 'ogg', 'webm', 'flv')


class ConfigError(Exception):
    pass


def get_videos_directory() -> pathlib.Path:
    """Get the "videos" directory in the media directory.  Make it if it does not exist."""
    directory = get_media_directory() / 'videos'
    if not directory.is_dir():
        directory.mkdir(parents=True)
    return directory


def get_no_channel_directory() -> pathlib.Path:
    """Get the "NO CHANNEL" directory in the videos directory.  Make it if it does not exist."""  # noqa
    directory = get_videos_directory() / 'NO CHANNEL'
    if not directory.is_dir():
        directory.mkdir(parents=True)
    return directory


def check_for_channel_conflicts(session: Session, id_=None, url=None, name=None, directory=None,
                                source_id=None):
    """
    Search for any channels that conflict with the provided args, raise a relevant exception if any conflicts are found.
    """
    if not any([id_, url, name, directory]):
        raise Exception('Cannot search for channel with no arguments')

    logger.debug(f'Checking for channel conflicts: id={id_} url={url} name={name} directory={directory}')

    # A channel can't conflict with itself
    if id_:
        base_where = session.query(Channel).filter(Channel.id != id_)
    else:
        base_where = session.query(Channel)

    if url:
        conflicts = base_where.filter(Channel.url == url)
        if list(conflicts):
            raise ChannelURLConflict()
    if name:
        conflicts = base_where.filter(Channel.name == name)
        if list(conflicts):
            raise ChannelNameConflict()
    if directory:
        conflicts = base_where.filter(Channel.directory == directory)
        if list(conflicts):
            raise ChannelDirectoryConflict()
    if source_id:
        conflicts = base_where.filter(Channel.source_id == source_id)
        if list(conflicts):
            raise ChannelSourceIdConflict()


def duration_to_seconds(duration: str) -> int:
    """Convert ffmpeg duration to seconds"""
    hours, minutes, seconds = duration.split(':')
    seconds, microseconds = seconds.split('.')
    return int(timedelta(
        hours=int(hours),
        minutes=int(minutes),
        seconds=int(seconds),
        microseconds=int(microseconds),
    ).total_seconds())


def seconds_to_duration(seconds: int) -> str:
    """Convert integer seconds to ffmpeg duration"""
    hours, seconds = divmod(seconds, 3600)
    minutes, seconds = divmod(seconds, 60)
    return f'{hours:02}:{minutes:02}:{seconds:02}.000'


DURATION_REGEX = re.compile(r'Duration: ([\d:.]+),')


def get_duration_from_ffmpeg(stderr: str) -> int:
    """Extract the duration from FFPEG's stderr, return this duration as integer seconds."""
    duration = DURATION_REGEX.findall(stderr)
    duration = duration_to_seconds(duration[0]) if duration else None
    return duration


def generate_video_poster(video_path: Path) -> Tuple[Path, Optional[int]]:
    """Create a poster (aka thumbnail) next to the provided video_path.  Also returns the video duration."""
    poster_path = video_path.with_suffix('.jpg')
    cmd = (FFMPEG_BIN, '-n', '-i', str(video_path), '-f', 'mjpeg', '-vframes', '1', '-ss', '00:00:05.000',
           str(poster_path))
    if poster_path.exists():
        # Poster already exists.
        return poster_path, None
    try:
        proc = subprocess.run(cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        stderr = proc.stderr.decode()
        logger.info(f'Generated poster at {poster_path}')
    except subprocess.CalledProcessError as e:
        logger.warning(f'FFMPEG poster generation failed with stdout: {e.stdout.decode()}')
        logger.warning(f'FFMPEG poster generation failed with stdout: {e.stderr.decode()}')
        raise
    if not poster_path.exists():
        raise Exception(f'Failed to find generated poster! {poster_path}')

    duration = get_duration_from_ffmpeg(stderr)
    return poster_path, duration


def convert_image(existing_path: Path, destination_path: Path, ext: str = 'jpeg'):
    """Convert an image from one format to another.  Remove the existing image file.  This will safely overwrite an
    image if the existing path is the same as the destination path.
    """
    with tempfile.NamedTemporaryFile(dir=destination_path.parent, delete=False) as fh:
        img = Image.open(existing_path).convert('RGB')
        img.save(fh.name, ext)

        existing_path.unlink()
        os.rename(fh.name, destination_path)
        os.chmod(destination_path, DEFAULT_FILE_PERMISSIONS)


def is_valid_poster(poster_path: Path) -> bool:
    """Return True only if poster file exists, and is a JPEG format."""
    if poster_path.is_file():
        try:
            img = Image.open(poster_path)
            return img.format == 'JPEG'
        except PIL.UnidentifiedImageError:
            logger.error(f'Failed to identify poster: {poster_path}', exc_info=True)
            pass

    return False


def get_video_duration(video_path: Path) -> int:
    """Get the duration of a video in seconds.  Do this using ffprobe."""
    if not isinstance(video_path, Path):
        video_path = Path(video_path)
    if not video_path.is_file():
        raise FileNotFoundError(f'{video_path} does not exist!')
    if not FFPROBE_BIN:
        raise SystemError('ffprobe is not installed!')

    cmd = [FFPROBE_BIN, '-v', 'error', '-show_entries', 'format=duration', '-of',
           'default=noprint_wrappers=1:nokey=1', str(video_path)]

    try:
        proc = subprocess.run(cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    except subprocess.CalledProcessError as e:
        logger.warning(f'FFPROBE failed to get duration with stdout: {e.stdout.decode()}')
        logger.warning(f'FFPROBE failed to get duration with stderr: {e.stderr.decode()}')
        raise
    stdout = proc.stdout.decode()
    duration = int(Decimal(stdout.strip()))
    return duration


def check_for_video_corruption(video_path: Path) -> bool:
    """Uses ffprobe to check for specific ways a video file can be corrupt.

    Also attempts to screenshot the video near the end of it's duration, this check if the video is corrupted somewhere
    in the middle."""
    video_path = Path(video_path) if not isinstance(video_path, Path) else video_path

    if not video_path.is_file():
        raise FileNotFoundError(f'{video_path} does not exist!')
    if not FFPROBE_BIN:
        raise SystemError('ffprobe is not installed!')

    # Just read the video using FFMPEG, this prints out useful information about the video.
    cmd = (FFPROBE_BIN, str(video_path))
    try:
        proc = subprocess.run(cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    except subprocess.CalledProcessError as e:
        logger.warning(f'FFPROBE failed to check for corruption with stderr: {e.stderr.decode()}')
        return True  # video is corrupt.

    messages = (
        b'Invalid NAL unit size',
        b'Error splitting the input into NAL units',
        b'Invalid data found when processing input',
    )
    corrupt = False
    for error in messages:
        if error in proc.stderr:
            logger.warning(f'Possible video corruption ({error.decode()}): {video_path}')
            corrupt = True
    return corrupt


def apply_info_json(channel_id: int):
    """Update view_count for all Videos in a channel using its info_json file.

    Mark any videos not in the info_json as "censored".
    """
    with get_db_session() as session:
        channel: Channel = session.query(Channel).filter_by(id=channel_id).one()
        channel_name = channel.name
        info = channel.info_json

    if not info:
        logger.info(f'No info_json for channel {channel.name}')
        return

    view_counts = [{'id': i['id'], 'view_count': i['view_count']} for i in info['entries']]
    view_counts_str = json.dumps(view_counts)

    with get_db_curs(commit=True) as curs:
        # Update the view_count for each video.
        stmt = '''
            WITH source AS (select * from json_to_recordset(%s::json) as (id text, view_count int))
            UPDATE video
            SET view_count = s.view_count
            FROM source as s
            WHERE source_id=s.id AND channel_id=%s
            RETURNING video.id AS updated_ids
        '''
        curs.execute(stmt, (view_counts_str, channel_id))
        count = len(curs.fetchall())
        logger.debug(f'Updated {count} view counts in DB for {channel_name}.')

        # Mark any video not in the info_json entries as censored.
        source_ids = [i['id'] for i in info['entries']]
        stmt = '''
            UPDATE video SET censored=(source_id != ALL(%s))
            WHERE channel_id=%s
        '''
        curs.execute(stmt, (source_ids, channel_id))


minimize_channel = partial(minimize_dict, keys=MINIMUM_CHANNEL_KEYS)
minimize_video_info_json = partial(minimize_dict, keys=MINIMUM_INFO_JSON_KEYS)
_minimize_video = partial(minimize_dict, keys=MINIMUM_VIDEO_KEYS)


def minimize_video(video: dict) -> dict:
    """
    Return a Video dictionary that contains only the key/values typically used.  Minimize the Channel and info_json,
    if they are present.
    """
    video = _minimize_video(video)

    if video.get('channel'):
        video['channel'] = minimize_channel(video['channel'])
    if video.get('info_json'):
        video['info_json'] = minimize_video_info_json(video['info_json'])

    return video


match_video_paths = partial(match_paths_to_suffixes, suffix_groups=(
    tuple(f'.{i}' for i in VIDEO_EXTENSIONS),
    ('.jpg', '.jpeg', '.webp', '.png'),
    ('.description',),
    ('.en.vtt', '.en.srt'),
    ('.info.json',),
))


@iterify(tuple)
def match_video_files(files: List[File]) -> Tuple[File, File, File, File, File]:
    video_paths = match_video_paths([i.path for i in files])
    for path in video_paths:
        yield next((i for i in files if i.path == path), None)