SaraView/Backend/AudioEmbed.py

"""Audio as a video, for link embeds (GET /public/embed/...).

Discord only plays audio that was uploaded to it: a link's og:audio is
ignored, and so is an og:video that points at an audio file. It does play an
MP4 og:video. So a shared audio file is turned into one with ffmpeg: a still
picture (the file's own cover art, or SaraView's icon) and the audio track.

A still frame encodes in seconds even for a long track. The video is made
once per file version (path, size and change time), kept in the data
directory's EmbedCache/, and started in the background as soon as an embed
page is built, so it's usually ready by the time someone presses play."""

from __future__ import annotations

import hashlib
import shutil
import subprocess
import threading
import tomllib
from pathlib import Path

from Settings import DefaultThemesDirectory, GetDataDirectory

ApplicationDirectory = Path(__file__).resolve().parent
# The browser-tab icon (Frontend/public/Icons), when the audio has no cover art.
IconPath = ApplicationDirectory.parent / "Frontend" / "dist" / "Icons" / "Icon.png"
FrameSize = 480
EncodeTimeout = 600  # seconds

_Locks: dict[str, threading.Lock] = {}
_LocksGuard = threading.Lock()


def IsAvailable() -> bool:
    return shutil.which("ffmpeg") is not None


_EncoderArguments: list[str] | None = None


def _VideoEncoder() -> list[str]:
    """libx264 where ffmpeg has it (Debian's, in the Docker image), else
    OpenH264 (Fedora's ffmpeg-free has only that). Both make H.264, which
    Discord plays."""
    global _EncoderArguments
    if _EncoderArguments is None:
        try:
            Encoders = subprocess.run(["ffmpeg", "-hide_banner", "-encoders"], capture_output=True, text=True,
                                      timeout=30).stdout
        except (OSError, subprocess.SubprocessError):
            Encoders = ""
        _EncoderArguments = (["-c:v", "libx264", "-tune", "stillimage", "-preset", "veryfast"] if " libx264 " in Encoders
                             else ["-c:v", "libopenh264"])
    return _EncoderArguments


def _BackgroundColour() -> str:
    """The default theme's Background, for the bars around non-square covers."""
    try:
        with (DefaultThemesDirectory / "CatppuccinMocha.toml").open("rb") as ThemeFile:
            return tomllib.load(ThemeFile)["Palette"]["Background"].lstrip("#")
    except (OSError, KeyError, tomllib.TOMLDecodeError):
        return "000000"


def _CachePath(AudioPath: Path) -> Path:
    Stat = AudioPath.stat()
    Key = hashlib.sha256(f"{AudioPath}|{Stat.st_size}|{Stat.st_mtime_ns}".encode()).hexdigest()[:32]
    return GetDataDirectory() / "EmbedCache" / f"{Key}.mp4"


def _ExtractCoverArt(AudioPath: Path, Target: Path) -> bool:
    """Writes the file's own cover art (if it has any) to Target as a PNG."""
    try:
        subprocess.run(["ffmpeg", "-y", "-loglevel", "error", "-i", str(AudioPath), "-an", "-frames:v", "1", str(Target)],
                       check=True, capture_output=True, timeout=60)
    except (OSError, subprocess.SubprocessError):
        return False
    return Target.is_file() and Target.stat().st_size > 0


def _Encode(AudioPath: Path, Target: Path) -> None:
    Target.parent.mkdir(parents=True, exist_ok=True)
    Partial = Target.with_suffix(".part.mp4")
    Cover = Target.with_suffix(".cover.png")
    Background = _BackgroundColour()
    Scaling = "neighbor"  # pixel art stays sharp
    if _ExtractCoverArt(AudioPath, Cover):
        Picture = ["-loop", "1", "-framerate", "1", "-i", str(Cover)]
        Scaling = "lanczos"  # album art is usually a photo
    elif IconPath.is_file():
        Picture = ["-loop", "1", "-framerate", "1", "-i", str(IconPath)]
    else:
        Picture = ["-f", "lavfi", "-i", f"color=c=0x{Background}:s={FrameSize}x{FrameSize}:r=1"]
    # A non-square cover gets bars in the theme's background colour.
    Fit = (f"scale={FrameSize}:{FrameSize}:force_original_aspect_ratio=decrease:flags={Scaling},"
           f"pad={FrameSize}:{FrameSize}:(ow-iw)/2:(oh-ih)/2:color=0x{Background},format=yuv420p")
    Command = [
        "ffmpeg", "-y", "-loglevel", "error", *Picture, "-i", str(AudioPath),
        "-map", "0:v:0", "-map", "1:a:0", "-vf", Fit, "-r", "1",
        *_VideoEncoder(),
        "-c:a", "aac", "-b:a", "192k", "-shortest", "-movflags", "+faststart", str(Partial),
    ]
    try:
        subprocess.run(Command, check=True, capture_output=True, timeout=EncodeTimeout)
        Partial.replace(Target)
    finally:
        Partial.unlink(missing_ok=True)
        Cover.unlink(missing_ok=True)


def GetEmbedVideo(AudioPath: Path) -> Path | None:
    """The embed video for an audio file, made now if needed; None if
    ffmpeg is missing or couldn't make it."""
    if not IsAvailable():
        return None
    Target = _CachePath(AudioPath)
    with _LocksGuard:
        Lock = _Locks.setdefault(str(Target), threading.Lock())
    with Lock:
        if not Target.is_file():
            try:
                _Encode(AudioPath, Target)
            except (OSError, subprocess.SubprocessError):
                return None
    return Target


def PrepareInBackground(AudioPath: Path) -> None:
    """Starts making the video while the embed page is being read."""
    if IsAvailable() and not _CachePath(AudioPath).is_file():
        threading.Thread(target=GetEmbedVideo, args=(AudioPath,), daemon=True).start()