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 pathlib import Path
from typing import Optional, Tuple, Union
from urllib.parse import urlparse

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

from wrolpi.captions import FFMPEG_BIN
from wrolpi.cmd import FFPROBE_BIN, run_command
from wrolpi.common import logger, get_media_directory
from wrolpi.dates import seconds_to_timestamp
from wrolpi.db import get_db_session, get_db_curs
from wrolpi.vars import DEFAULT_FILE_PERMISSIONS
from .errors import ChannelNameConflict, ChannelURLConflict, ChannelDirectoryConflict, ChannelSourceIdConflict

logger = logger.getChild(__name__)


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."""
    from modules.videos.lib import format_videos_destination
    videos_destination = format_videos_destination()
    directory = get_media_directory() / videos_destination
    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 get_youtube_channel_id(url: str) -> Optional[str]:
    """Extract YouTube channel ID (UC...) from /channel/ URLs."""
    parsed = urlparse(url)
    if parsed.netloc not in ('youtube.com', 'www.youtube.com', 'm.youtube.com'):
        return None
    if parsed.path.startswith('/channel/'):
        parts = parsed.path.split('/')
        if len(parts) >= 3 and parts[2]:
            return parts[2]
    return None


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.
    """
    from .models import Channel

    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:
        from wrolpi.collections import Collection
        conflicts = base_where.join(Collection).filter(Collection.name == name)
        if list(conflicts):
            raise ChannelNameConflict()
    if directory:
        from wrolpi.collections import Collection
        conflicts = base_where.join(Collection).filter(Collection.directory == directory)
        if list(conflicts):
            raise ChannelDirectoryConflict()
    if source_id:
        conflicts = base_where.filter(Channel.source_id == source_id)
        if list(conflicts):
            raise ChannelSourceIdConflict()

    # Assure that this Channel's directory is not in another Channel's directory.
    directory = f'{directory}/'
    for channel in session.query(Channel).filter(Channel.id != id_):
        channel_directory = f'{channel.directory}/'
        if directory.startswith(channel_directory) or channel_directory.startswith(directory):
            raise ChannelDirectoryConflict()


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 ffmpeg_poster(video_path: Path, poster_path: Path, seconds: int):
    """Use ffmpeg to create a poster file from a video file at a particular time of the video.

    Returns stderr so data about the command can be extracted."""
    timestamp = seconds_to_timestamp(seconds)
    cmd = (FFMPEG_BIN, '-ss', timestamp, '-n', '-i', video_path, '-f', 'mjpeg', '-vframes', '1', poster_path)
    try:
        proc = subprocess.run(cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        stderr = proc.stderr.decode()
    except subprocess.CalledProcessError as e:
        logger.error(f'Failed to create video poster', exc_info=e)
        raise

    return stderr


def ffmpeg_video_complete(video_path: Path, seconds: int = None) -> bool:
    """Checks if video file is complete by taking screenshot from the end of the video.

    @raise FileNotFoundError: raised if the video file does not exist.
    """
    if not video_path.is_file():
        raise FileNotFoundError(f'Video file not found: {video_path}')
    if not video_path.stat().st_size:
        raise RuntimeError(f'Video file is empty: {video_path}')

    with tempfile.NamedTemporaryFile(suffix='.jpg', delete=False) as fh:
        path = Path(fh.name)
        path.unlink()
        try:
            seconds = seconds or extract_video_duration(video_path) - 5
            ffmpeg_poster(video_path, path, seconds)
            return True
        except subprocess.CalledProcessError:
            return False
        except TypeError:
            # "N/A" duration.
            return False
        finally:
            if path.is_file():
                path.unlink()


def generate_video_poster(video_path: Path, seconds: int = 5) -> 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')
    if poster_path.exists():
        # Poster already exists.
        return poster_path, None

    stderr = ffmpeg_poster(video_path, poster_path, seconds)

    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 destination file already exists.
    """
    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


async def ffprobe_json(video_path: Union[Path, str]) -> dict:
    """Extract video file metadata using ffprobe.

    >>> ffprobe_json('video.mp4')
    {
        'chapters': [],
        'format': {},
        'streams': [
            {},
        ],
    }
    """
    if not FFPROBE_BIN:
        raise RuntimeError('ffprobe was not found')

    cmd = (FFPROBE_BIN,
           '-print_format', 'json',
           '-loglevel', 'quiet',
           '-show_streams',
           '-show_format',
           '-show_chapters',
           video_path.absolute())

    result = await run_command(cmd, timeout=60)

    if result.return_code != 0:
        raise RuntimeError(f'Got non-zero exit code: {result.return_code}')

    try:
        content = json.loads(result.stdout.decode().strip())
    except Exception as e:
        logger.debug(result.stdout.decode())
        logger.error(f'Failed to load ffprobe json', exc_info=e)
        raise

    if content == dict():
        # Data is empty, video may be corrupt.
        raise RuntimeError('ffprobe data was empty')

    return content


def ffprobe_json_sync(video_path: Union[Path, str]) -> dict:
    """Synchronous version of ffprobe_json for use in sync validation code.

    Extract video file metadata using ffprobe.

    >>> ffprobe_json_sync('video.mp4')
    {
        'chapters': [],
        'format': {},
        'streams': [
            {},
        ],
    }
    """
    if not FFPROBE_BIN:
        raise RuntimeError('ffprobe was not found')

    cmd = [FFPROBE_BIN,
           '-print_format', 'json',
           '-loglevel', 'quiet',
           '-show_streams',
           '-show_format',
           '-show_chapters',
           str(Path(video_path).absolute())]

    proc = subprocess.run(cmd, capture_output=True, timeout=60)

    if proc.returncode != 0:
        raise RuntimeError(f'Got non-zero exit code: {proc.returncode}')

    try:
        content = json.loads(proc.stdout.decode().strip())
    except Exception as e:
        logger.debug(proc.stdout.decode())
        logger.error(f'Failed to load ffprobe json', exc_info=e)
        raise

    if content == dict():
        # Data is empty, video may be corrupt.
        raise RuntimeError('ffprobe data was empty')

    return content


def read_ffprobe_json_file(path: Path) -> Optional[dict]:
    """Read ffprobe data from .ffprobe.json file if it exists."""
    if not path or not path.is_file():
        return None
    try:
        with open(path, 'r') as f:
            return json.load(f)
    except (json.JSONDecodeError, IOError) as e:
        logger.warning(f'Failed to read ffprobe json file {path}', exc_info=e)
        return None


def write_ffprobe_json_file(path: Path, data: dict):
    """Write ffprobe data to .ffprobe.json file with human-readable formatting."""
    with open(path, 'w') as f:
        json.dump(data, f, indent=2)


async def get_or_create_ffprobe_json(video_path: Path) -> Tuple[dict, Optional[Path]]:
    """Return `(ffprobe data, the cache file written)` for a video file.

    Deliberately takes a path and touches no session: ffprobe is a subprocess that runs for
    seconds, and a caller holding a write transaction while it runs blocks every other writer
    until `busy_timeout` expires.  The caller persists the returned data afterwards, in a short
    write transaction of its own.
    """
    ffprobe_path = video_path.with_suffix('.ffprobe.json')
    if ffprobe_path.is_file() and (cached_data := read_ffprobe_json_file(ffprobe_path)):
        return cached_data, None

    data = await ffprobe_json(video_path)

    written = None
    from wrolpi.common import get_wrolpi_config
    if get_wrolpi_config().save_ffprobe_json:
        try:
            write_ffprobe_json_file(ffprobe_path, data)
            written = ffprobe_path
        except IOError as e:
            logger.warning(f'Failed to write ffprobe json cache file {ffprobe_path}', exc_info=e)
    return data, written


def extract_video_duration(video_path: Path) -> Optional[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().strip()
    if stdout == 'N/A':
        return None
    duration = int(Decimal(stdout))
    return duration


async def update_view_counts_and_censored(channel_id: int):
    """Update view_count for all Videos in a channel using its info_json file.  Also sets FileGroup.censored
    if Video is no longer available on the Channel."""
    from .models import Channel

    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'] if i.get('view_count') is not None]
    view_counts_str = json.dumps(view_counts)

    with get_db_curs(commit=True) as curs:
        # Update the view_count for each video.
        stmt = '''
               UPDATE video
               SET view_count = CAST(json_extract(s.value, '$.view_count') AS INTEGER)
               FROM json_each(:view_counts) AS s
               WHERE video.source_id = json_extract(s.value, '$.id')
                 AND video.channel_id = :channel_id
               RETURNING video.id \
               '''
        curs.execute(stmt, {'view_counts': view_counts_str, 'channel_id': channel_id})
        count = len(curs.fetchall())
        logger.info(f'Updated {count} view counts in DB for {channel_name}.')

    source_ids = [i['id'] for i in info['entries']]
    with get_db_curs(commit=True) as curs:
        # Set FileGroup.censored if the video is no longer on the Channel.
        stmt = '''
               UPDATE file_group AS fg
               SET censored = (v.source_id NOT IN (SELECT value FROM json_each(:source_ids)))
               FROM video v
               WHERE v.file_group_id = fg.id
                 AND v.channel_id = :channel_id
               RETURNING id, censored \
               '''
        curs.execute(stmt, {'channel_id': channel_id, 'source_ids': json.dumps(source_ids)})
        censored = len([i for i in curs.fetchall() if i['censored']])
        logger.info(f'Set {censored} censored videos for {channel_name}.')