voice-cat/clients/apple/dotnet/VoiceCat.Mac/ScreenAudioCapture.cs

using System.Runtime.InteropServices;
using CoreFoundation;
using CoreMedia;
using Foundation;
using ObjCRuntime;
using ScreenCaptureKit;
using VoiceCat.Audio;

namespace VoiceCat.Mac;

internal enum ScreenAudioScope { EntireDesktop, OnlyApplications, AllExceptApplications }
internal sealed record ScreenAudioSelection(ScreenAudioScope Scope, IReadOnlySet<string> BundleIdentifiers, bool ExcludeScreenReader)
{
    internal static ScreenAudioSelection Default { get; } = new(ScreenAudioScope.EntireDesktop, new HashSet<string>(), false);
}

internal sealed class ScreenAudioCapture : NSObject, ISCStreamOutput
{
    private static readonly HashSet<string> ScreenReaders = ["com.apple.VoiceOver", "com.apple.VoiceOver4", "com.apple.speech.speechsynthesisd"];
    private readonly int channels;
    private readonly CapturePcmHandler onPcm;
    private readonly ScreenAudioSelection selection;
    private readonly DispatchQueue queue = new("net.iamtalon.voicecat.screen-audio");
    private readonly CaptureDelegate streamDelegate;
    private SCStream? stream;
    private int disposed;

    internal ScreenAudioCapture(int channels, ScreenAudioSelection selection, CapturePcmHandler onPcm)
    {
        this.channels = Math.Clamp(channels, 1, 2); this.selection = selection; this.onPcm = onPcm;
        streamDelegate = new(error => Failed?.Invoke(error));
    }

    internal event Action<Exception>? Failed;

    internal static async Task<IReadOnlyList<ScreenApplication>> GetApplicationsAsync(CancellationToken cancellationToken = default)
    {
        SCShareableContent content = await GetContentAsync(cancellationToken).ConfigureAwait(false);
        return content.Applications.Where(application => !string.IsNullOrWhiteSpace(application.BundleIdentifier) &&
                !ScreenReaders.Contains(application.BundleIdentifier) && application.BundleIdentifier != NSBundle.MainBundle.BundleIdentifier)
            .GroupBy(application => application.BundleIdentifier, StringComparer.Ordinal).Select(group => group.First())
            .Select(application => new ScreenApplication(string.IsNullOrWhiteSpace(application.ApplicationName) ? application.BundleIdentifier : application.ApplicationName,
                application.BundleIdentifier)).OrderBy(application => application.Name, StringComparer.CurrentCultureIgnoreCase).ToArray();
    }

    internal async Task StartAsync(CancellationToken cancellationToken = default)
    {
        SCShareableContent content = await GetContentAsync(cancellationToken).ConfigureAwait(false);
        SCDisplay display = content.Displays.FirstOrDefault() ?? throw new InvalidOperationException("No display is available for screen audio capture.");
        SCRunningApplication[] matching = content.Applications.Where(application => Included(application.BundleIdentifier)).ToArray();
        SCContentFilterOption option = selection.Scope == ScreenAudioScope.OnlyApplications ? SCContentFilterOption.Include : SCContentFilterOption.Exclude;
        var filter = new SCContentFilter(display, matching, [], option);
        var configuration = new SCStreamConfiguration
        {
            CapturesAudio = true, ExcludesCurrentProcessAudio = true, SampleRate = 48_000,
            ChannelCount = channels, Width = 2, Height = 2, QueueDepth = 6,
            MinimumFrameInterval = new CMTime(1, 1)
        };
        var created = new SCStream(filter, configuration, streamDelegate);
        if (!created.AddStreamOutput(this, SCStreamOutputType.Audio, queue, out NSError? outputError))
            throw new InvalidOperationException(outputError?.LocalizedDescription ?? "Could not attach the screen-audio output.");
        stream = created;
        var started = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
        created.StartCapture(error => { if (error is null) started.TrySetResult(); else started.TrySetException(new InvalidOperationException(error.LocalizedDescription)); });
        await started.Task.WaitAsync(TimeSpan.FromSeconds(15), cancellationToken).ConfigureAwait(false);
    }

    private bool Included(string bundleIdentifier)
    {
        bool selected = selection.BundleIdentifiers.Contains(bundleIdentifier);
        bool reader = selection.ExcludeScreenReader && ScreenReaders.Contains(bundleIdentifier);
        return selection.Scope switch
        {
            ScreenAudioScope.OnlyApplications => selected,
            ScreenAudioScope.AllExceptApplications => selected || reader,
            _ => reader
        };
    }

    [Export("stream:didOutputSampleBuffer:ofType:")]
    public unsafe void DidOutputSampleBuffer(SCStream captureStream, CMSampleBuffer sampleBuffer, SCStreamOutputType type)
    {
        if (type != SCStreamOutputType.Audio || !sampleBuffer.DataIsReady || disposed != 0) return;
        int frames = checked((int)sampleBuffer.NumSamples);
        if (frames <= 0 || frames > 8192) return;
        Span<byte> listStorage = stackalloc byte[40];
        fixed (byte* list = listStorage)
        {
            int status = CMSampleBufferGetAudioBufferListWithRetainedBlockBuffer(sampleBuffer.Handle, out _, (IntPtr)list,
                (nuint)listStorage.Length, IntPtr.Zero, IntPtr.Zero, 0, out IntPtr blockBuffer);
            if (status != 0) return;
            try
            {
                uint count = *(uint*)list;
                if (count is 0 or > 2) return;
                Span<short> converted = stackalloc short[frames * channels];
                NativeAudioBuffer* first = (NativeAudioBuffer*)(list + 8);
                bool planar = count > 1;
                int sourceChannels = planar ? checked((int)count) : checked((int)Math.Max(1, first->Channels));
                for (int frame = 0; frame < frames; frame++)
                {
                    float left = planar ? ((float*)first[0].Data)[frame] : ((float*)first[0].Data)[frame * sourceChannels];
                    float right = sourceChannels > 1
                        ? planar ? ((float*)first[1].Data)[frame] : ((float*)first[0].Data)[frame * sourceChannels + 1]
                        : left;
                    if (channels == 1) converted[frame] = ToInt16(sourceChannels > 1 ? (left + right) * 0.5f : left);
                    else { converted[frame * 2] = ToInt16(left); converted[frame * 2 + 1] = ToInt16(right); }
                }
                onPcm(converted, channels);
            }
            finally { if (blockBuffer != IntPtr.Zero) CFRelease(blockBuffer); }
        }
    }

    private static short ToInt16(float value) => (short)Math.Clamp((int)MathF.Round(Math.Clamp(value, -1, 1) * 32767), short.MinValue, short.MaxValue);

    private static Task<SCShareableContent> GetContentAsync(CancellationToken cancellationToken)
    {
        var completion = new TaskCompletionSource<SCShareableContent>(TaskCreationOptions.RunContinuationsAsynchronously);
        SCShareableContent.GetShareableContent((content, error) =>
        {
            if (error is not null) completion.TrySetException(new UnauthorizedAccessException(error.LocalizedDescription));
            else if (content is null) completion.TrySetException(new InvalidOperationException("ScreenCaptureKit returned no shareable content."));
            else completion.TrySetResult(content);
        });
        cancellationToken.Register(() => completion.TrySetCanceled(cancellationToken));
        return completion.Task;
    }

    protected override void Dispose(bool disposing)
    {
        if (disposing && Interlocked.Exchange(ref disposed, 1) == 0)
        {
            SCStream? previous = Interlocked.Exchange(ref stream, null);
            if (previous is null) { streamDelegate.Dispose(); queue.Dispose(); }
            else previous.StopCapture(_ => { previous.Dispose(); streamDelegate.Dispose(); queue.Dispose(); });
        }
        base.Dispose(disposing);
    }

    [StructLayout(LayoutKind.Sequential)]
    private struct NativeAudioBuffer { internal uint Channels, DataByteSize; internal IntPtr Data; }

    private sealed class CaptureDelegate(Action<Exception> failed) : SCStreamDelegate
    {
        public override void DidStop(SCStream stream, NSError error) => failed(new IOException(error.LocalizedDescription));
    }

    [DllImport("/System/Library/Frameworks/CoreMedia.framework/CoreMedia")]
    private static extern int CMSampleBufferGetAudioBufferListWithRetainedBlockBuffer(IntPtr sampleBuffer,
        out nuint bufferListSizeNeeded, IntPtr bufferList, nuint bufferListSize, IntPtr structureAllocator,
        IntPtr blockAllocator, uint flags, out IntPtr blockBuffer);

    [DllImport("/System/Library/Frameworks/CoreFoundation.framework/CoreFoundation")]
    private static extern void CFRelease(IntPtr value);
}