bernini_continuation.py

from __future__ import annotations

import base64
import re
import subprocess
import tempfile
from pathlib import Path
from typing import Any


_DEADLINE = None

VIDEO_SUFFIXES = {'.mp4', '.mov', '.mkv', '.webm', '.avi', '.m4v'}


def _safe_name(value: object, fallback: str) -> str:
    name = str(value or fallback).replace('\\', '/').split('/')[-1]
    name = re.sub(r'[^A-Za-z0-9._-]+', '_', name).strip('._')
    return (name or fallback)[:160]


def _decode(value: object, label: str) -> bytes:
    encoded = str(value or '').strip()
    if encoded.startswith('data:') and ',' in encoded:
        encoded = encoded.split(',', 1)[1]
    encoded = re.sub(r'\s+', '', encoded).replace('-', '+').replace('_', '/')
    encoded += '=' * ((4 - len(encoded) % 4) % 4)
    try:
        data = base64.b64decode(encoded, validate=True)
    except Exception as exc:
        raise ValueError(f'{label} is not valid base64') from exc
    if not data:
        raise ValueError(f'{label} is empty')
    return data


def _run(command: list[str], label: str, timeout: int) -> None:
    if _DEADLINE is not None: timeout = _DEADLINE.remaining(timeout)
    completed = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, timeout=timeout, check=False)
    if completed.returncode:
        detail = (completed.stderr or completed.stdout or '').strip()[-2000:]
        raise RuntimeError(f'{label} failed: {detail}')


def validate_continuation(value: object) -> dict[str, Any] | None:
    if value is None or value is False:
        return None
    if not isinstance(value, dict):
        raise ValueError('continuation must be an object')
    if value.get('enabled') is False:
        return None
    source_name = _safe_name(value.get('source_video_name'), 'bernini_source_video.mp4')
    frame_name = _safe_name(value.get('last_frame_name'), 'bernini_continuation_last_frame.png')
    if Path(source_name).suffix.lower() not in VIDEO_SUFFIXES:
        raise ValueError('continuation source must be a video file')
    if Path(frame_name).suffix.lower() != '.png':
        raise ValueError('continuation last frame must be a PNG file')
    try:
        width = int(value.get('width'))
        height = int(value.get('height'))
        fps = int(value.get('fps', 20))
        drop = int(value.get('drop_generated_frames', 2))
        target_value = value.get('target_duration_seconds')
        target_duration = float(target_value) if target_value is not None else None
    except (TypeError, ValueError) as exc:
        raise ValueError('invalid continuation settings') from exc
    if not 256 <= width <= 2048 or not 256 <= height <= 2048 or width % 2 or height % 2:
        raise ValueError('invalid continuation dimensions')
    if not 1 <= fps <= 60 or not 1 <= drop <= 10:
        raise ValueError('invalid continuation frame settings')
    if target_duration is not None and not 1 <= target_duration <= 3600:
        raise ValueError('invalid continuation target duration')
    return {'source_video_name': source_name, 'last_frame_name': frame_name, 'width': width, 'height': height, 'fps': fps, 'drop_generated_frames': drop, 'target_duration_seconds': target_duration}


def prepare_continuation_assets(payload: dict[str, Any]) -> dict[str, Any] | None:
    config = validate_continuation(payload.get('continuation'))
    if config is None:
        return None
    assets = payload.get('assets')
    if not isinstance(assets, list):
        raise ValueError('continuation assets are missing')
    source = next((item for item in assets if isinstance(item, dict) and item.get('name') == config['source_video_name']), None)
    if source is None:
        raise ValueError('continuation source video was not supplied')
    source_bytes = _decode(source.get('data'), config['source_video_name'])
    suffix = Path(config['source_video_name']).suffix.lower()
    with tempfile.TemporaryDirectory(prefix='video-continuation-frame-') as temp_dir:
        root = Path(temp_dir)
        source_path = root / ('source' + suffix)
        output_path = root / config['last_frame_name']
        source_path.write_bytes(source_bytes)
        resize = f"scale={config['width']}:{config['height']}:force_original_aspect_ratio=decrease,pad={config['width']}:{config['height']}:(ow-iw)/2:(oh-ih)/2:color=black,setsar=1"
        errors = []
        for offset in ('-0.1', '-0.5', '-1.0'):
            output_path.unlink(missing_ok=True)
            try:
                _run(
                    ['ffmpeg', '-hide_banner', '-loglevel', 'error', '-y', '-sseof', offset,
                     '-i', str(source_path), '-map', '0:v:0', '-an', '-frames:v', '1',
                     '-vf', resize, str(output_path)],
                    'continuation last-frame extraction', 300,
                )
            except (RuntimeError, subprocess.TimeoutExpired) as exc:
                errors.append(str(exc))
            if output_path.is_file() and output_path.stat().st_size > 0:
                break
        else:
            output_path.unlink(missing_ok=True)
            try:
                _run(
                    ['ffmpeg', '-hide_banner', '-loglevel', 'error', '-y', '-i', str(source_path),
                     '-map', '0:v:0', '-an', '-vf', resize, '-fps_mode', 'passthrough',
                     '-update', '1', str(output_path)],
                    'continuation full-decode fallback', 3600,
                )
            except (RuntimeError, subprocess.TimeoutExpired) as exc:
                errors.append(str(exc))
        if not output_path.is_file() or output_path.stat().st_size <= 0:
            raise RuntimeError('continuation last-frame extraction produced no image: ' + '; '.join(errors)[-2000:])
        assets.append({'name': config['last_frame_name'], 'data': base64.b64encode(output_path.read_bytes()).decode('ascii')})
    return config


def merge_continuation_output(source_path: Path, generated_path: Path, output_path: Path, config: dict[str, Any]) -> None:
    output_path.parent.mkdir(parents=True, exist_ok=True)
    resize = f"scale={config['width']}:{config['height']}:force_original_aspect_ratio=decrease,pad={config['width']}:{config['height']}:(ow-iw)/2:(oh-ih)/2:color=black,setsar=1,format=yuv420p"
    filters = (
        f"[0:v]fps={config['fps']},{resize},setpts=PTS-STARTPTS[source];"
        f"[1:v]fps={config['fps']},{resize},trim=start_frame={config['drop_generated_frames']},setpts=PTS-STARTPTS[generated];"
        '[source][generated]concat=n=2:v=1:a=0[video]'
    )
    partial = output_path.with_suffix(output_path.suffix + '.part.mp4')
    _run([
        'ffmpeg', '-hide_banner', '-loglevel', 'error', '-y', '-i', str(source_path), '-i', str(generated_path),
        '-filter_complex', filters, '-map', '[video]', '-map', '0:a?', '-c:v', 'libx264', '-preset', 'fast',
        '-crf', '16', '-pix_fmt', 'yuv420p', '-c:a', 'aac', '-b:a', '192k',
        *(['-t', str(config['target_duration_seconds'])] if config['target_duration_seconds'] is not None else []),
        '-movflags', '+faststart', str(partial),
    ], 'continuation video merge', 3600)
    if not partial.is_file() or partial.stat().st_size <= 0:
        raise RuntimeError('continuation merge produced no video')
    partial.replace(output_path)