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
# Part of every cache key: bump it when the videos change, so old ones
# aren't served (2: cut to the audio's exact length).
EncodingVersion = 2
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"{EncodingVersion}|{AudioPath}|{Stat.st_size}|{Stat.st_mtime_ns}".encode()).hexdigest()[:32]
    return GetDataDirectory() / "EmbedCache" / f"{Key}.mp4"


def _AudioDuration(AudioPath: Path) -> float | None:
    try:
        Result = subprocess.run(["ffprobe", "-v", "error", "-show_entries", "format=duration", "-of", "csv=p=0",
                                 str(AudioPath)], capture_output=True, text=True, timeout=60)
        return float(Result.stdout.strip())
    except (OSError, subprocess.SubprocessError, ValueError):
        return None


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")
    # The picture loops forever, so the video is cut to the audio's length.
    # -shortest alone overshoots: libx264 buffers ~40 frames ahead, which at
    # one frame a second left the picture running 40+ seconds past the
    # audio, and Discord wouldn't play that.
    Duration = _AudioDuration(AudioPath)
    Length = ["-t", f"{Duration:.3f}"] if Duration else []
    Command = [
        "ffmpeg", "-y", "-loglevel", "error", *Picture, "-i", str(AudioPath),
        "-map", "0:v:0", "-map", "1:a:0", "-vf", Fit, "-r", "1", "-g", "10",
        *_VideoEncoder(),
        "-c:a", "aac", "-b:a", "192k", "-shortest", *Length, "-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()