koibot/src/utils/sessions.ts

import { Collection } from "@discordjs/collection"
import { db, gateway, voiceStateManager } from "../api.js"
import {
	AudioReceiveStream,
	DiscordGatewayAdapterLibraryMethods,
	EndBehaviorType,
	joinVoiceChannel,
	VoiceConnection,
} from "@discordjs/voice"
import { Snowflake } from "discord-api-types/v10"
import * as prism from "prism-media"
import { pipeline } from "stream"
import { createWriteStream, existsSync, mkdirSync } from "fs"
import path from "path"
import { isMe } from "./appid.js"
import { exec, spawn } from "child_process"

export class RecordingSessionManager {
	sessionMap = new Collection<Snowflake, RecordingSession>()
	adapters = new Collection<Snowflake, DiscordGatewayAdapterLibraryMethods>()

	hasSession(guild: Snowflake): boolean {
		return this.sessionMap.has(guild)
	}

	async addSession(
		guild: Snowflake,
		voiceChannel: Snowflake,
		textChannel: Snowflake,
		shardId: number
	) {
		const session = new RecordingSession(
			joinVoiceChannel({
				guildId: guild,
				channelId: voiceChannel,
				adapterCreator: (methods) => {
					this.adapters.set(guild, methods)
					return {
						sendPayload: (payload) => {
							gateway.send(shardId, payload)
							return true
						},
						destroy: () => {
							this.adapters.delete(guild)
						},
					}
				},
				selfDeaf: false,
				selfMute: true,
				debug: true,
			}),
			voiceChannel,
			textChannel
		)

		await session.init()

		this.sessionMap.set(guild, session)
	}

	removeSession(guild: Snowflake) {
		const session = this.sessionMap.get(guild)
		if (session) session.shutdown()
		this.sessionMap.delete(guild)
	}
}

export class RecordingSession {
	streams = new Collection<Snowflake, UserStream>()
	startedAt = Date.now()
	id!: string

	constructor(
		public vc: VoiceConnection,
		public voiceChannelId: Snowflake,
		public textChannelId: Snowflake
	) {}

	async init() {
		const result = await db.session.create({
			data: {
				id: this.id,
				channel: this.voiceChannelId,
			},
		})
		this.id = result.id

		this.id = Date.now().toString()
		this.receiver.speaking.on("start", (userId) => this.addReceiver(userId))
		this.receiver.speaking.on("end", (userId) =>
			this.removeReceiver(userId)
		)
	}

	get receiver() {
		return this.vc.receiver
	}

	addReceiver(userId: string) {
		if (isMe(userId)) return
		if (this.streams.has(userId)) return

		const receiver = this.receiver.subscribe(userId)

		this.streams.set(userId, new UserStream(this, receiver, userId))
	}

	removeReceiver(userId: string) {
		const stream = this.streams.get(userId)

		if (stream) {
			stream.subscription.push(null)
			this.streams.delete(userId)
		}
	}

	async shutdown() {}
}

class UserStream {
	oggStream: prism.opus.OggLogicalBitstream

	constructor(
		public session: RecordingSession,
		public subscription: AudioReceiveStream,
		public userId: string
	) {
		this.oggStream = new prism.opus.OggLogicalBitstream({
			opusHead: new prism.opus.OpusHead({
				channelCount: 2,
				sampleRate: 48000,
			}),
			pageSizeControl: {
				maxPackets: 10,
			},
		})

		const dir = path.resolve("recordings", session.id)
		const timestamp = Date.now() - session.startedAt
		const file = `${userId}-${timestamp}.ogg`

		if (!existsSync(dir)) {
			mkdirSync(dir, {
				recursive: true,
			})
		}

		const filename = path.join(dir, file)

		console.log(`Recording ${filename} started`)

		const outputStream = createWriteStream(filename)

		pipeline(subscription, this.oggStream, outputStream, async (err) => {
			outputStream.end()

			if (err) {
				console.log(`Recording ${filename} ended with error: ${err}`)
				subscription.destroy(err)
				return
			}

			await db.file
				.create({
					data: {
						session: {
							connect: {
								id: session.id,
							},
						},
						filename: file,
						timestamp,
					},
				})
				.catch(console.error)

			console.log(`Recording ${filename} ended`)
		})
	}
}