discord-meeting-bot/recording.py

"""Recording sinks for meeting capture.

Hybrid capture strategy:

* one mixed archive WAV (everyone blended, silence-padded for a faithful
  timeline) - good for archival/playback
* one speech-only WAV per Discord speaker - good for speaker-attributed
  transcription
* a ``tracks.json`` manifest mapping each per-speaker file back to the Discord
  user that produced it

The mixed archive keeps the existing behaviour (wrap with the library's
``SilenceGeneratorSink`` so dropped packets/gaps don't collapse the timeline).
The per-speaker files intentionally skip synthetic silence frames so they stay
small and contain only that person's speech.
"""

from __future__ import annotations

import json
import logging
import wave
from pathlib import Path
from typing import Optional

from discord.ext import voice_recv

logger = logging.getLogger(__name__)

# Discord voice receive decodes to 48kHz, 16-bit, stereo PCM.
CHANNELS = 2
SAMPLE_WIDTH = 2
SAMPLING_RATE = 48000

MANIFEST_NAME = "tracks.json"


def _has_audio(pcm: Optional[bytes]) -> bool:
    """True when pcm contains at least one non-zero byte.

    Synthetic silence frames from the silence generator are all-zero, so this
    lets us keep per-speaker tracks limited to real speech.
    """
    if not pcm:
        return False
    return pcm.count(0) != len(pcm)


def _display_name(user) -> str:
    if user is None:
        return "Unknown User"
    name = getattr(user, "display_name", None) or getattr(user, "name", None)
    if name:
        return str(name)
    user_id = getattr(user, "id", None)
    return f"User {user_id}" if user_id is not None else "Unknown User"


class MeetingSink(voice_recv.AudioSink):
    """Writes a mixed archive WAV plus one speech-only WAV per speaker."""

    def __init__(self, mixed_path: str, tracks_dir: str):
        super().__init__()
        self.mixed_path = mixed_path
        self.tracks_dir = Path(tracks_dir)
        self._cleaned_up = False

        Path(self.mixed_path).parent.mkdir(parents=True, exist_ok=True)
        self.tracks_dir.mkdir(parents=True, exist_ok=True)

        self._mixed = self._open_wave(self.mixed_path)
        # user_id -> {"writer": Wave_write, "path": str, "display_name": str}
        self._tracks: dict[int, dict] = {}

    @staticmethod
    def _open_wave(path: str) -> wave.Wave_write:
        writer = wave.open(path, "wb")
        writer.setnchannels(CHANNELS)
        writer.setsampwidth(SAMPLE_WIDTH)
        writer.setframerate(SAMPLING_RATE)
        return writer

    def wants_opus(self) -> bool:
        return False

    def _track_for(self, user) -> Optional[dict]:
        user_id = getattr(user, "id", None)
        if user_id is None:
            return None

        track = self._tracks.get(user_id)
        if track is None:
            path = self.tracks_dir / f"track-{user_id}.wav"
            track = {
                "writer": self._open_wave(str(path)),
                "path": str(path),
                "display_name": _display_name(user),
            }
            self._tracks[user_id] = track
            logger.info(
                "Started per-speaker track: user_id=%s display_name=%s path=%s",
                user_id,
                track["display_name"],
                path,
            )
        else:
            # Display name can resolve later than the first packet.
            resolved = _display_name(user)
            if resolved != "Unknown User":
                track["display_name"] = resolved
        return track

    def write(self, user, data) -> None:
        pcm = getattr(data, "pcm", b"") or b""

        # Mixed archive always gets every frame (including silence) so the
        # timeline stays faithful.
        try:
            self._mixed.writeframes(pcm)
        except Exception:
            logger.exception("Failed writing to mixed archive: path=%s", self.mixed_path)

        # Per-speaker track only gets real speech.
        if user is None or not _has_audio(pcm):
            return

        track = self._track_for(user)
        if track is None:
            return
        try:
            track["writer"].writeframes(pcm)
        except Exception:
            logger.exception("Failed writing per-speaker track: path=%s", track["path"])

    def _write_manifest(self) -> None:
        manifest = {
            "mixed_path": self.mixed_path,
            "tracks": [
                {
                    "user_id": user_id,
                    "display_name": track["display_name"],
                    "path": track["path"],
                }
                for user_id, track in sorted(self._tracks.items())
            ],
        }
        manifest_path = self.tracks_dir / MANIFEST_NAME
        manifest_path.write_text(json.dumps(manifest, indent=2))
        logger.info(
            "Wrote speaker-track manifest: path=%s speakers=%s",
            manifest_path,
            len(manifest["tracks"]),
        )

    def cleanup(self) -> None:
        if self._cleaned_up:
            return
        self._cleaned_up = True

        # Close per-speaker writers first, then the manifest, then the mixed
        # archive LAST. That ordering guarantees: if the mixed WAV is a valid,
        # finalized file, the manifest already exists on disk.
        for user_id, track in self._tracks.items():
            try:
                track["writer"].close()
            except Exception:
                logger.warning("Error closing per-speaker track user_id=%s", user_id, exc_info=True)

        try:
            self._write_manifest()
        except Exception:
            logger.exception("Failed writing speaker-track manifest")

        try:
            self._mixed.close()
        except Exception:
            logger.warning("Error closing mixed archive on cleanup", exc_info=True)