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


# How much of a video's tail is decoded to prove the file is complete.  Corruption from an
# interrupted download or write lands at the end of the file; decoding all of it would take as
# long as playing it.
TAIL_DECODE_SECONDS = 30
TAIL_DECODE_TIMEOUT = 10 * 60


def tail_decode_command(video_path: Path, seconds: int = None) -> Tuple[str | Path, ...]:
    """ffmpeg command which decodes the last `seconds` of `video_path` to nothing, printing only
    errors.  `-sseof` seeks from the end; a file shorter than `seconds` is decoded whole."""
    seconds = seconds or TAIL_DECODE_SECONDS
    return FFMPEG_BIN, '-v', 'error', '-sseof', f'-{int(seconds)}', '-i', video_path, '-f', 'null', '-'


def _tail_decode_verdict(return_code: int, stderr: bytes | str) -> Tuple[bool, str]:
    errors = (stderr.decode(errors='replace') if isinstance(stderr, bytes) else (stderr or '')).strip()
    ok = return_code == 0 and not errors
    return ok, errors


def _check_video_file(video_path: Path):
    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}')
    if not FFMPEG_BIN:
        raise RuntimeError('ffmpeg was not found')


def ffmpeg_video_complete(video_path: Path, seconds: int = None) -> bool:
    """True if the last `seconds` of the video decode without error.

    Used after a download, and after a transcode before the original is replaced.  A truncated or
    corrupt file fails here because its damage is at the end; the index (moov) must also be
    readable for ffmpeg to seek from the end at all.

    @raise FileNotFoundError: raised if the video file does not exist.
    @raise RuntimeError: raised if the video file is empty, or ffmpeg is missing.
    """
    _check_video_file(video_path)
    cmd = tuple(str(i) for i in tail_decode_command(video_path, seconds))
    try:
        proc = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, timeout=TAIL_DECODE_TIMEOUT)
    except subprocess.TimeoutExpired:
        logger.error(f'Timed out checking video is complete: {video_path}')
        return False
    ok, errors = _tail_decode_verdict(proc.returncode, proc.stderr)
    if not ok:
        logger.warning(f'Video is incomplete or corrupt: {video_path} (exit {proc.returncode}) {errors[-500:]}')
    return ok


async def ffmpeg_video_complete_async(video_path: Path, seconds: int = None, runner=None) -> Tuple[bool, str]:
    """`ffmpeg_video_complete` for async callers; returns (complete, ffmpeg's error output).

    `runner` defaults to `run_command`; a Job passes its own so ffmpeg's output lands in the Job log.
    """
    _check_video_file(video_path)
    runner = runner or run_command
    result = await runner(tail_decode_command(video_path, seconds), cwd=video_path.parent, timeout=TAIL_DECODE_TIMEOUT)
    ok, errors = _tail_decode_verdict(result.return_code, result.stderr)
    if not ok:
        logger.warning(f'Video is incomplete or corrupt: {video_path} (exit {result.return_code}) {errors[-500:]}')
    return ok, errors


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}.')