nonsensebot/app/modules/ts_transcode.py

"""
This module actively listens for messages in a specific room from specific users, and checks if they contain
a video which uses an unsupported video codec. If it does, it will be transcoded and re-uploaded to the room.

This is used because we mirror Truth Social into a matrix room and sometimes the videos are in H265, which a lot of
clients cannot play.

This can easily be adapted to other rooms and users, or even to other types of media, by changing the USER_IDS list.
"""

import logging
import shlex
import shutil
import typing
import urllib.parse

import niobot
import tempfile
import asyncio
from pathlib import Path


class TruthSocialTranscode(niobot.Module):
    USER_IDS = {"@trumpbot:shronk.net", "@tatebot:shronk.net", "@nex:nexy7574.co.uk"}

    def __init__(self, bot: niobot.NioBot):
        super().__init__(bot)
        self.log = logging.getLogger(__name__)
        if not shutil.which("ffprobe"):
            raise RuntimeError("ffprobe is required for this module.")
        if not shutil.which("ffmpeg"):
            raise RuntimeError("ffmpeg is required for this module.")

        # noinspection PyTypeChecker
        self.bot.add_event_callback(self.on_message, niobot.RoomMessageVideo)

    def __teardown__(self):
        for callback in self.bot.event_callbacks:
            if callback.func == self.on_message:
                self.bot.event_callbacks.remove(callback)
        super().__teardown__()

    async def on_message_internal(
        self, room: niobot.MatrixRoom, event: niobot.RoomMessageVideo
    ) -> str:
        if self.bot.is_old(event):
            return "IGNORED"                            
        self.log.info("Processing event: %r (%r)", event, room)
        # if not isinstance(event, niobot.RoomMessageVideo):
        #     # We are not interested in non-videos.
        #     return
        # event: niobot.RoomMessageVideo

        if event.sender not in self.USER_IDS:
            # We are not interested in messages from other users.
            self.log.info("Ignored user: %s", event.sender)
            return "IGNORED"

        # Videos sent in the truth room are not encrypted, so that makes this easier.
        http_url = await self.bot.mxc_to_http(
            event.url,
            homeserver=event.sender.split(":", 1)[
                1
            ],  # use the sender homeserver for download
        )
        http_url = "https://" + http_url
        name = event.flattened()["content.filename"]

        with tempfile.NamedTemporaryFile(suffix="-%s" % name) as temp_fd:
            # async with self.bot.client_session.get(http_url) as resp:
            #     temp_fd.write(await resp.read())
            #     temp_fd.flush()
            dl = await self.bot.download(event.url, save_to=Path(temp_fd.name))
            if isinstance(dl, niobot.DownloadError):
                self.log.error("Error downloading video: %r", dl)
                await self.bot.send_message(
                    room,
                    "Error downloading video, unable to transcode: %r" % dl,
                    reply_to=event.event_id,
                )
                return "ERROR_DOWNLOAD"

            # First, probe the video to see if it has H265 content in it.
            probe = await asyncio.create_subprocess_exec(
                "ffprobe",
                "-v",
                "error",
                "-select_streams",
                "v:0",
                "-show_entries",
                "stream=codec_name",
                "-of",
                "default=noprint_wrappers=1:nokey=1",
                temp_fd.name,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
            )
            stdout, stderr = await probe.communicate()
            if probe.returncode != 0:
                # Something went wrong, let's not continue.
                self.log.error(
                    "Error probing video:\n%s", (stderr or b"EOF").decode("utf-8")
                )
                await self.bot.send_message(
                    room,
                    "Error probing video, unable to determine whether a transcode is required.",
                    reply_to=event.event_id,
                )
                return "ERROR_PROBE"

            codec = stdout.decode("utf-8").strip()
            if codec in ["h264", "vp8", "vp9"]:
                # We don't need to transcode this.
                self.log.info("Video is already in a supported format: %s", codec)
                return "NO_NEED"

            # We now need to transcode this. For compatibility, we'll force it to be H264 and AAC.
            # We will, however, use the High profile for H264, as almost everything supports it now, plus better
            # compression.
            # We will also limit the bitrate to about 5 megabit, because Trump has a habit of posting raw footage
            # that ends up being something stupid like 20 megabit at 60fps.
            with tempfile.NamedTemporaryFile(
                suffix="-transcoded-" + name + ".mp4"
            ) as out_fd:
                self.log.info("transcoding %s -> %s", temp_fd.name, out_fd.name)
                args = [
                    "ffmpeg",
                    "-v",
                    "warning",
                    "-i",
                    temp_fd.name,
                    "-c:v",
                    "libx264",
                    "-pix_fmt",
                    "yuv420p",
                    "-profile:v",
                    "high",
                    "-preset",
                    "medium",
                    "-b:v",
                    "5M",
                    "-maxrate",
                    "5M",
                    "-bufsize",
                    "10M",
                    "-minrate",
                    "1M",
                    "-c:a",
                    "aac",
                    "-b:a",
                    "96k",
                    "-vf",
                    "scale=-2:'min(1080,ih)',fps=25",
                    "-movflags",
                    "faststart",
                    "-y",
                    # "-fps_mode",
                    # "vfr",
                    "-f",
                    "mp4",
                    out_fd.name,
                ]
                transcode = await asyncio.create_subprocess_exec(
                    *args,
                    stdout=asyncio.subprocess.PIPE,
                    stderr=asyncio.subprocess.PIPE,
                )
                self.log.info("Starting transcode...")
                self.log.debug("Running command: %r", shlex.join(args))
                stdout, stderr = await transcode.communicate()
                self.log.info("Finished transcode.")
                if transcode.returncode != 0:
                    self.log.error(
                        "Error transcoding video:\n%s\n\n%s",
                        stderr.decode("utf-8"),
                        stdout.decode("utf-8"),
                    )
                    await self.bot.send_message(
                        room,
                        "Error transcoding video, unable to transcode: got return code %d"
                        % transcode.returncode,
                        reply_to=event.event_id,
                    )
                    return "ERROR_TRANSCODE"
                # We transcoded! Now we need to upload it.
                file = await niobot.VideoAttachment.from_file(out_fd.name)
                self.log.info("Uploading transcoded video")
                await self.bot.send_message(
                    room.room_id, file=file, reply_to=event.event_id
                )
                self.log.info("Uploaded transcoded video")
                return "OK"

    async def on_message(self, room, event):
        try:
            return await self.on_message_internal(room, event)
        except Exception as e:
            self.log.exception("Error processing video", exc_info=e)
            return "ERROR %r" % e

    @staticmethod
    def is_allowed(ctx: niobot.Context) -> bool:
        return ctx.message.sender.split(":", 1)[1] in {
            "shronk.net",
            "nicroxio.co.uk",
            "nexy7574.co.uk",
        }

    @niobot.command("transcode")
    @niobot.from_homeserver(
        "nexy7574.co.uk", "shronk.net", "nicroxio.co.uk", "transgender.ing"
    )
    async def do_transcode(
        self,
        ctx: niobot.Context,
        message_link: typing.Annotated[
            niobot.RoomMessage, niobot.EventParser("m.room.message")
        ],
    ):
        """
        Transcodes a video to H264/AAC.
        """
        if not isinstance(message_link, niobot.RoomMessageVideo):
            await ctx.respond("That's not a video. %r" % message_link)
            return
        res = await ctx.respond("\N{HOURGLASS} Dispatching, please wait...")
        ret = await self.on_message(ctx.room, message_link)
        await res.edit(
            f"\N{WHITE HEAVY CHECK MARK} Transcode complete. Result: `{ret}`"
        )