voice-cat/clients/apple/VoiceCat.iOS/IosAudioEngine.cs

using System.Runtime.InteropServices;
using AVFoundation;
using UIKit;
using VoiceCat.Audio;
using VoiceCat.Core;

namespace VoiceCat.iOS;

internal sealed class IosAudioEngine
{
    // InstallTapOnBus bufferSize is only a request. Physical iOS hardware can deliver 4,800
    // frames per callback (observed with the voice-processing graph), so conversion storage
    // must cover substantially more than the requested 960 frames without allocating in Capture.
    private const int MaximumCaptureCallbackFrames = 16_384;
    internal static IosAudioEngine Shared { get; } = new();
    private readonly AdaptivePcmBuffer playbackRing = new(2, capacityFrames: 65_536);
    private readonly short[] renderScratch = new short[16_384];
    private AVAudioEngine? engine;
    private Foundation.NSObject? engineConfigurationObserver;
    private AVAudioSourceNode? source;
    private AVAudioFormat? outputFormat;
    private VoiceCatClient? client;
    private sealed class MicrophoneRoute(uint streamId, int channels)
    {
        internal readonly uint StreamId = streamId;
        internal readonly int Channels = channels;
        // Capture hardware and the managed 20 ms sender have independent clocks. Correct
        // their small rate difference before the queue eventually reaches its hard edge.
        // RemoteIO can deliver capture in 100 ms bursts, so retain one burst of headroom.
        internal readonly AdaptivePcmBuffer Ring = new(channels, 120, 65_536);
    }
    private MicrophoneRoute? microphone;
    private readonly short[] microphoneFrame = new short[960 * 2];
    private int microphoneCredit;
    private AVAudioFormat? microphoneFormat;
    private AVAudioConverter? microphoneConverter;
    private AVAudioPcmBuffer? convertedMicrophone;
    private AVAudioPcmBuffer? pendingInput;
    private AVAudioConverterInputHandler? inputProvider;
    private bool inputProvided;
    private bool tapInstalled;
    // The voice-processing state the live graph was actually built with, so Reconfigure can tell
    // a settings change apart from a re-check of an unchanged graph.
    private bool voiceProcessing;
    private long captureCallbacks, capturedFrames, convertedFrames, rejectedFeeds, converterFailures, stereoFrames, stereoDifferentFrames;
    private long renderCallbacks, lastRenderTimestamp;
    internal bool IsConnected { get; private set; }
    internal bool IsRunning => engine?.Running == true;
    // The watchdog compares this across ticks: a graph that stops calling back while it still
    // reports Running leaves the whole device-clocked pipeline frozen until it is rebuilt.
    internal long RenderCallbacks => Interlocked.Read(ref renderCallbacks);
    internal int BufferMilliseconds { get => playbackRing.BufferMilliseconds; set => playbackRing.BufferMilliseconds = value; }

    internal void StartListening(VoiceCatClient owner)
    {
        if (client is { } previous) previous.Audio.MixedPcm -= ReceiveMixedPcm;
        client = owner; IsConnected = true; owner.Audio.MixedPcm += ReceiveMixedPcm;
        // A reconnect after Detach finds the session active and the graph already running on the
        // right hardware. Rebuilding it there would release an HFP headset and pay a Bluetooth
        // profile renegotiation for a transport blip that changed no audio configuration.
        if (engine?.Running == true) { playbackRing.Resynchronize(); return; }
        Volatile.Write(ref microphone, null); Rebuild();
    }

    // Unbinds the client without touching the session or the graph, for a connection that was
    // lost rather than ended. Capture keeps feeding a ring nobody drains and Render emits silence
    // until StartListening rebinds, which keeps the route and its Bluetooth profile alive.
    internal void Detach()
    {
        if (client is { } owner) owner.Audio.MixedPcm -= ReceiveMixedPcm;
        client = null; playbackRing.Resynchronize();
        // The route's stream id belongs to the connection that just died. Keep the tap and its
        // hardware, but park the route on the unbound id so a rebind cannot feed the next
        // connection a stream it never announced.
        if (Volatile.Read(ref microphone) is { } stale && stale.StreamId != 0)
            Volatile.Write(ref microphone, CreateMicrophoneRoute(0, stale.Channels));
    }

    internal void StartMicrophone(uint streamId, int channels)
    {
        MicrophoneRoute? current = Volatile.Read(ref microphone);
        // Restoring voice after a reconnect only changes which stream id the existing capture
        // feeds. Reuse a running tap of the same width instead of rebuilding the graph around it.
        if (current is { StreamId: 0 } && current.Channels == Math.Clamp(channels, 1, 2) && tapInstalled && engine?.Running == true)
        {
            current.Ring.Resynchronize();
            Volatile.Write(ref microphone, CreateMicrophoneRoute(streamId, channels)); return;
        }
        Volatile.Write(ref microphone, CreateMicrophoneRoute(streamId, channels)); Rebuild();
    }

    internal void StopMicrophone() { Volatile.Write(ref microphone, null); Rebuild(); }

    // `force` rebuilds unconditionally, which is what a route change, a media-services reset and
    // the stall watchdog all need. Callers that are only re-checking a graph they expect to be
    // healthy — foregrounding, above all — pass false and get a no-op when nothing has changed.
    internal void Reconfigure(bool force = true)
    {
        if (!IsConnected) return;
        MicrophoneRoute? current = Volatile.Read(ref microphone);
        int channels = IosAudioRouter.Shared.CaptureChannels;
        if (!force && engine?.Running == true && tapInstalled == (current is not null) &&
            (current?.Channels ?? channels) == channels && voiceProcessing == IosAudioRouter.Shared.UsesVoiceProcessing) return;
        if (current is not null)
        {
            // Stop the old tap before publishing a route with a different sample width. Otherwise
            // an in-flight callback could interpret its old converter buffer using the new width.
            DestroyGraph();
            if (current.Channels != channels && current.StreamId != 0) client?.Audio.SetCaptureChannels(current.StreamId, channels);
            Volatile.Write(ref microphone, CreateMicrophoneRoute(current.StreamId, channels));
        }
        Rebuild();
    }

    private static MicrophoneRoute CreateMicrophoneRoute(uint streamId, int channels)
    {
        return new MicrophoneRoute(streamId, Math.Clamp(channels, 1, 2));
    }
    internal bool EnsureRunning()
    {
        if (!IsConnected || engine?.Running == true) return true;
        Rebuild(); return engine?.Running == true;
    }

    private void Rebuild()
    {
        DestroyGraph(); MicrophoneRoute? route = Volatile.Read(ref microphone); IosAudioRouter.Shared.Apply(route is not null);
        voiceProcessing = IosAudioRouter.Shared.UsesVoiceProcessing;
        var next = new AVAudioEngine();
        AVAudioInputNode? input = null;
        if (route is not null)
        {
            // Enabling voice processing rebuilds both sides of AVAudioEngine. Do it before any
            // formats are queried or nodes are connected so the graph is built from the final IO.
            input = next.InputNode;
            if (!input.SetVoiceProcessingEnabled(IosAudioRouter.Shared.UsesVoiceProcessing, out NSError? processingError))
                throw new InvalidOperationException(processingError.LocalizedDescription);
            if (IosAudioRouter.Shared.UsesVoiceProcessing) input.VoiceProcessingAgcEnabled = IosAudioRouter.Shared.AutomaticGainControl;
        }
        outputFormat = new(AVAudioCommonFormat.PCMFloat32, 48_000, 2, false);
        source = new(outputFormat, Render);
        next.AttachNode(source);
        NSError? connectionError = null;
        if (OperatingSystem.IsIOSVersionAtLeast(27)) next.Connect(source, next.MainMixerNode, outputFormat, out connectionError);
        else next.Connect(source, next.MainMixerNode, outputFormat);
        if (connectionError is not null) throw new InvalidOperationException(connectionError.LocalizedDescription);
        if (route is not null)
        {
            AVAudioFormat inputFormat = input!.GetBusOutputFormat(0);
            Console.Error.WriteLine($"VC_GRAPH vpio={IosAudioRouter.Shared.UsesVoiceProcessing} requestedCh={route.Channels} " +
                $"inputCh={inputFormat.ChannelCount} inputRate={inputFormat.SampleRate}");
            microphoneFormat = new(AVAudioCommonFormat.PCMInt16, 48_000, (uint)route.Channels, true);
            microphoneConverter = new(inputFormat, microphoneFormat);
            uint capacity = checked((uint)Math.Ceiling(MaximumCaptureCallbackFrames * 48_000 / inputFormat.SampleRate) + 64);
            convertedMicrophone = new(microphoneFormat, capacity);
            inputProvider = ProvideInput;
            NSError? tapError = null;
            if (OperatingSystem.IsIOSVersionAtLeast(27)) input.InstallTapOnBus(0, 960, inputFormat, out tapError, Capture);
            else input.InstallTapOnBus(0, 960, inputFormat, Capture);
            if (tapError is not null) throw new InvalidOperationException(tapError.LocalizedDescription);
            tapInstalled = true;
        }
        next.Prepare();
        if (!next.StartAndReturnError(out NSError? error)) { next.Dispose(); throw new InvalidOperationException(error.LocalizedDescription); }
        engine = next;
        engineConfigurationObserver = Foundation.NSNotificationCenter.DefaultCenter.AddObserver(
            AVAudioEngine.ConfigurationChangeNotification, notification =>
            {
                if (!ReferenceEquals(notification.Object, next)) return;
                UIApplication.SharedApplication.BeginInvokeOnMainThread(() =>
                {
                    if (IsConnected && ReferenceEquals(engine, next) && !next.Running) Rebuild();
                });
            }, next);
    }

    private unsafe void Capture(AVAudioPcmBuffer buffer, AVAudioTime time)
    {
        Interlocked.Increment(ref captureCallbacks);
        Interlocked.Add(ref capturedFrames, buffer.FrameLength);
        VoiceCatClient? owner = client; MicrophoneRoute? route = Volatile.Read(ref microphone);
        if (owner is null || route is null || buffer.FrameLength == 0) return;
        AVAudioConverter? converter = microphoneConverter;
        AVAudioPcmBuffer? converted = convertedMicrophone;
        AVAudioConverterInputHandler? provider = inputProvider;
        if (converter is null || converted is null || provider is null) return;
        pendingInput = buffer; inputProvided = false; converted.FrameLength = 0;
        converter.ConvertToBuffer(converted, out NSError? conversionError, provider);
        if (conversionError is not null) Interlocked.Increment(ref converterFailures);
        if (converted.FrameLength == 0) return;
        Interlocked.Add(ref convertedFrames, converted.FrameLength);
        nint samples = Marshal.ReadIntPtr(converted.Int16ChannelData);
        if (samples != 0 && ReferenceEquals(route, Volatile.Read(ref microphone)))
        {
            var pcm = new ReadOnlySpan<short>((void*)samples, checked((int)converted.FrameLength * route.Channels));
            if (route.Channels == 2)
            {
                Interlocked.Add(ref stereoFrames, converted.FrameLength);
                long different = 0;
                for (int frame = 0; frame < converted.FrameLength; frame++)
                    if (pcm[frame * 2] != pcm[frame * 2 + 1]) different++;
                Interlocked.Add(ref stereoDifferentFrames, different);
            }
            if (!route.Ring.TryWrite(pcm)) Interlocked.Increment(ref rejectedFeeds);
        }
        pendingInput = null;
    }

    internal string CaptureDiagnostics()
    {
        return $"callbacks={Interlocked.Exchange(ref captureCallbacks, 0)} input={Interlocked.Exchange(ref capturedFrames, 0)} " +
            $"converted={Interlocked.Exchange(ref convertedFrames, 0)} feedDrops={Interlocked.Exchange(ref rejectedFeeds, 0)} " +
            $"converterErrors={Interlocked.Exchange(ref converterFailures, 0)} stereo={Interlocked.Exchange(ref stereoDifferentFrames, 0)}/" +
            $"{Interlocked.Exchange(ref stereoFrames, 0)}";
    }

    private AVAudioBuffer ProvideInput(uint _, out AVAudioConverterInputStatus status)
    {
        if (!inputProvided && pendingInput is { } input) { inputProvided = true; status = AVAudioConverterInputStatus.HaveData; return input; }
        status = AVAudioConverterInputStatus.NoDataNow; return null!;
    }

    // One paced 20 ms handoff from the capture ring into the encoder. The render callback calls
    // this at its demand rate: the ring absorbs RemoteIO's burst pattern and the small capture vs
    // output clock difference, so the sender sees a steady 20 ms feed without a sleep-paced pacer
    // thread to stall when iOS coalesces a backgrounded app's wakeups.
    private void PumpMicrophoneChunk(VoiceCatClient owner)
    {
        MicrophoneRoute? route = Volatile.Read(ref microphone);
        // Stream id 0 is a route parked by Detach: still capturing, not yet bound to a connection.
        if (route is null || route.StreamId == 0) return;
        int required = 960 * route.Channels;
        if (route.Ring.Read(microphoneFrame.AsSpan(0, required)) == required &&
            ReferenceEquals(route, Volatile.Read(ref microphone)) &&
            !owner.Audio.FeedPcm(route.StreamId, microphoneFrame.AsSpan(0, required), route.Channels))
            Interlocked.Increment(ref rejectedFeeds);
    }

    private void ReceiveMixedPcm(ReadOnlySpan<short> pcm) => playbackRing.TryWrite(pcm);

    private unsafe int Render(IntPtr isSilence, IntPtr timestamp, uint frameCount, IntPtr outputData)
    {
        int frames = checked((int)frameCount), requested = checked(frames * 2);
        if (requested > renderScratch.Length) return -1;
        // This callback is the cadence iOS keeps exact while the app is backgrounded or the device
        // is locked, so it owns both managed 20 ms hands-offs: capture into the sender and one mix
        // cycle per 20 ms of render demand. Both run allocation-free and without locks or I/O.
        Interlocked.Increment(ref renderCallbacks);
        long now = System.Diagnostics.Stopwatch.GetTimestamp(), previous = lastRenderTimestamp;
        lastRenderTimestamp = now;
        VoiceCatClient? owner = Volatile.Read(ref client);
        if (owner is null || !IsConnected) microphoneCredit = 0;
        else
        {
            if (previous != 0 && (now - previous) * 1000.0 / System.Diagnostics.Stopwatch.Frequency > 100)
            {
                Volatile.Read(ref microphone)?.Ring.Resynchronize();
                playbackRing.Resynchronize(); owner.Audio.ResynchronizeInputs(); microphoneCredit = 0;
            }
            microphoneCredit += frames;
            while (microphoneCredit >= 960) { microphoneCredit -= 960; PumpMicrophoneChunk(owner); }
            // Top the mix ring up to cover this callback plus its configured target so the read
            // below never starves and the buffer keeps its chosen buffering latency. Catch-up is
            // capped at one extra cycle so a refill cannot overrun this callback's deadline.
            int deficit = Math.Min(frames + playbackRing.TargetFrames - playbackRing.CountFrames, frames + 960);
            for (int produced = 0; produced < deficit; produced += 960) owner.Audio.RunCycle();
        }
        Span<short> input = renderScratch.AsSpan(0, requested);
        int read = playbackRing.Read(input); input[read..].Clear();
        int count = Marshal.ReadInt32(outputData), first = IntPtr.Size == 8 ? 8 : 4, stride = IntPtr.Size == 8 ? 16 : 12;
        if (count != 2) return -1;
        for (int channel = 0; channel < 2; channel++)
        {
            nint data = Marshal.ReadIntPtr(outputData, first + channel * stride + 8);
            var output = new Span<float>((void*)data, frames);
            for (int frame = 0; frame < frames; frame++) output[frame] = input[frame * 2 + channel] / 32768f;
        }
        if (isSilence != IntPtr.Zero) Marshal.WriteByte(isSilence, read == 0 ? (byte)1 : (byte)0);
        return 0;
    }

    internal void Stop()
    {
        IsConnected = false; Volatile.Write(ref microphone, null);
        if (client is { } owner) owner.Audio.MixedPcm -= ReceiveMixedPcm;
        DestroyGraph(); client = null; IosAudioRouter.Shared.Deactivate();
    }

    private void DestroyGraph()
    {
        if (engineConfigurationObserver is { } observer)
        {
            Foundation.NSNotificationCenter.DefaultCenter.RemoveObserver(observer);
            observer.Dispose(); engineConfigurationObserver = null;
        }
        if (engine is { } old)
        {
            if (tapInstalled) old.InputNode.RemoveTapOnBus(0);
            old.Stop(); if (source is not null) old.DetachNode(source); old.Dispose();
        }
        tapInstalled = false; pendingInput = null; inputProvider = null; lastRenderTimestamp = 0; microphoneCredit = 0;
        convertedMicrophone?.Dispose(); convertedMicrophone = null;
        microphoneConverter?.Dispose(); microphoneConverter = null;
        microphoneFormat?.Dispose(); microphoneFormat = null;
        source?.Dispose(); source = null; outputFormat?.Dispose(); outputFormat = null; engine = null;
    }
}