voice-cat/clients/apple/iOS/VoiceCatiOS/IOSVoiceProcessingEngine.swift

import AVFoundation
import Darwin
import os
import VoiceCatCore

private let logger = Logger(subsystem: "cat.voice.VoiceCatiOS", category: "IOSAudioEngine")

/// In-process single-producer/single-consumer int16 PCM ring for the playback path.
///
///   producer = the core's mixer-timer thread (the `vc_set_mixed_output_sink` callback)
///   consumer = the `AVAudioSourceNode` render thread
///
/// Heap-backed (not shared memory like `BroadcastAudioRing`), but the same discipline: aligned
/// 64-bit monotonic indices with `OSMemoryBarrier` for acquire/release ordering. Both the C
/// callback and the render block are real-time — they only do index math + a memcpy here, never
/// lock or allocate.
final class PCMRing {
    private let data: UnsafeMutablePointer<Int16>
    private let capacity: Int
    private var writeIdx: UInt64 = 0
    private var readIdx: UInt64 = 0

    init(capacitySamples: Int) {
        capacity = capacitySamples
        data = UnsafeMutablePointer<Int16>.allocate(capacity: capacitySamples)
        data.initialize(repeating: 0, count: capacitySamples)
    }
    deinit { data.deallocate() }

    /// Producer: append `count` interleaved int16 samples. Drops the chunk if it doesn't fit
    /// (better to skip than tear). Single producer only (the core mixer-timer thread).
    func write(_ src: UnsafePointer<Int16>, count: Int) {
        guard count > 0, count <= capacity else { return }
        let w = writeIdx
        OSMemoryBarrier()
        let r = readIdx
        if capacity - Int(w &- r) < count { return }  // full: drop
        var idx = Int(w % UInt64(capacity))
        var off = 0
        var rem = count
        while rem > 0 {
            let chunk = min(rem, capacity - idx)
            (data + idx).update(from: src + off, count: chunk)
            idx = (idx + chunk) % capacity
            off += chunk
            rem -= chunk
        }
        OSMemoryBarrier()
        writeIdx = w &+ UInt64(count)
    }

    /// Consumer: read up to `count` interleaved int16 samples into `dst`; returns the number
    /// read (the rest is the caller's to silence-fill). Single consumer only (render thread).
    func read(into dst: UnsafeMutablePointer<Int16>, count: Int) -> Int {
        let r = readIdx
        OSMemoryBarrier()
        let w = writeIdx
        let available = Int(w &- r)
        if available <= 0 { return 0 }
        let n = min(available, count)
        var idx = Int(r % UInt64(capacity))
        var off = 0
        var rem = n
        while rem > 0 {
            let chunk = min(rem, capacity - idx)
            (dst + off).update(from: data + idx, count: chunk)
            idx = (idx + chunk) % capacity
            off += chunk
            rem -= chunk
        }
        OSMemoryBarrier()
        readIdx = r &+ UInt64(n)
        return n
    }

    /// Consumer-side snapshot of how many interleaved int16 samples are currently buffered. Lets a
    /// paced consumer check for a full frame *before* calling `read`, so it never reads (and thus
    /// discards) a partial frame. Single consumer only (same thread that calls `read`).
    var availableSamples: Int {
        let r = readIdx
        OSMemoryBarrier()
        let w = writeIdx
        return Int(w &- r)
    }

    /// Discard everything buffered — call before (re)starting so stale pre-roll isn't played.
    func reset() { OSMemoryBarrier(); readIdx = writeIdx }

    /// Diagnostics: monotonic total samples written / read since the ring was created. The
    /// indices are already cumulative, so these are free. Only read them when both threads are
    /// quiesced (e.g. at teardown after the engine + mixer sink are stopped) — they are not
    /// synchronized for live cross-thread reads. Lets us tell "core never delivered PCM" apart
    /// from "PCM arrived but produced no sound" (the AVAudioEngine output graph).
    var debugTotalWritten: UInt64 { writeIdx }
    var debugTotalRead: UInt64 { readIdx }
}

/// The single iOS audio engine (docs/voice.md §8 "iOS audio engine").
///
/// **One path, always external.** On iOS the core never opens a miniaudio device: a MIC stream is
/// always started with `external_feed=1`, `vc_set_external_playback(1)` is set once at connect, and
/// this engine drives *both* directions through one `AVAudioEngine`:
///   - **core → speaker:** the core's mixed-output sink fills `ring`; an `AVAudioSourceNode` pulls
///     from it and renders through the engine output. This runs the whole time we're connected,
///     so remote audio plays even before the user joins voice (no "can't hear anyone").
///   - **mic → core:** when the mic is active a tap on the input node converts to 48 kHz int16 and
///     writes to a pacing ring; a 20 ms timer releases steady 960-sample frames to
///     `client.feedPcm(micStreamId)`. The core sends each captured frame synchronously, so this
///     steady cadence is what keeps packets from bursting and fluttering the receiver's playout.
///
/// Echo cancellation / noise suppression / AGC come from Apple's Voice-Processing I/O unit (VPIO),
/// which `inputNode.setVoiceProcessingEnabled(true)` enables. VPIO forces mono, so it is engaged
/// only when the active preset wants it (`IOSAudioRouter.currentConfigUsesVoiceProcessing`) — the
/// Stereo Mic / A2DP configs run the same engine with VPIO off.
///
/// Every preset / route / interruption change funnels through `reconfigure()`: a single
/// deterministic stop → AVAudioSession reconfigure → rebuild graph → start. There is no second
/// (miniaudio) audio path to hand off to, so a switch cannot leave one direction dropped.
@MainActor
final class IOSAudioEngine {
    static let shared = IOSAudioEngine()

    /// True while connected (between `startListening` and `stop`) — the playback graph should run.
    private(set) var isConnected = false
    /// True while a local mic stream is active — the input tap should be installed.
    private(set) var micActive = false

    private let engine = AVAudioEngine()
    private var sourceNode: AVAudioSourceNode?
    private weak var client: VoiceCatClient?
    private var micStreamId: UInt32 = 0
    private var captureChannels: UInt32 = 1

    // AVAudioEngine may deliver several codec frames per callback. Pace complete 20 ms frames
    // through an SPSC ring; never consume partial frames, and recreate the timer when the channel
    // count changes.
    private let micRing = PCMRing(capacitySamples: 48000 * 2)  // ~1 s stereo — ample elastic slack
    private var micTimer: DispatchSourceTimer?
    private let micQueue = DispatchQueue(label: "cat.voice.mic.feedPump")
    private let micDrainScratch: UnsafeMutablePointer<Int16>
    private static let micFrameSamplesPerChannel = 960  // 20 ms @ 48 kHz — core's frame size

    /// Feed-pump state, touched only on `micQueue` (the pump's serial queue). A reference type so
    /// the timer closure mutates it without capturing `self` (which is @MainActor). `targetFrames`
    /// is the prebuffer depth: the pump fills this many frames before it starts releasing, so the
    /// tap's bursty delivery (~2 frames at once) can't drain it to empty between bursts. It persists
    /// across rebuilds and self-heals upward (capped) on an underrun, so it tunes to whatever IO
    /// buffer size the active route/VPIO actually uses without a hard-coded guess.
    private final class PumpState {
        var primed = false
        var targetFrames = 3  // ~60 ms initial cushion; grows on underrun up to maxTargetFrames
        static let maxTargetFrames = 6  // ~120 ms cap — bounds added latency
    }
    private let pumpState = PumpState()

    // 48 kHz stereo Float32 (deinterleaved) — the format the source node renders. The core
    // delivers 48 kHz stereo int16 via the mixed-output sink; mainMixerNode adapts to the route.
    private let outFormat = AVAudioFormat(
        commonFormat: .pcmFormatFloat32, sampleRate: 48000, channels: 2, interleaved: false)!

    // Playback ring (mixed remote audio): ~0.5 s of 48 kHz stereo int16. Filled by the core's
    // mixer-timer thread, drained by the source-node render thread.
    private let ring = PCMRing(capacitySamples: 48000 * 2 / 2)
    // Render-thread scratch for deinterleaving — pre-allocated so the render block never allocates.
    private let renderScratchFrames = 8192
    private let renderScratch: UnsafeMutablePointer<Int16>

    private init() {
        renderScratch = UnsafeMutablePointer<Int16>.allocate(capacity: renderScratchFrames * 2)
        renderScratch.initialize(repeating: 0, count: renderScratchFrames * 2)
        micDrainScratch = UnsafeMutablePointer<Int16>.allocate(capacity: 960 * 2)
        micDrainScratch.initialize(repeating: 0, count: 960 * 2)

        // AVAudioEngine stops itself on a mid-session route/configuration change (it stops
        // if its I/O graph no longer matches the active route). Our route-change handler in
        // AudioSessionManager normally rebuilds us before the user notices, but if the engine
        // stops itself AFTER our recovery (because the route-change notification raced ahead
        // of the engine's own self-stop), nothing restarts it. Catch that case here.
        NotificationCenter.default.addObserver(
            self, selector: #selector(handleEngineConfigurationChange),
            name: .AVAudioEngineConfigurationChange, object: engine)
    }

    /// The engine stopped itself because its configuration no longer matches the active AVAudio
    /// route (this fires after a route change that the route-change handler can't always outrun).
    /// Dispatch to main and call the unified `recoverAudio()` — it's intent-gated on
    /// `isConnected`, idempotent, and no-ops if the engine is already running (the common case
    /// where our route-change handler got there first).
    @objc private func handleEngineConfigurationChange(_ notification: Notification) {
        Task { @MainActor [weak self] in
            guard let self else { return }
            guard self.isConnected, !self.engine.isRunning else { return }
            logger.info("engine configuration-change — engine stopped itself, recovering")
            AudioSessionManager.shared.recoverAudio()
        }
    }

    // MARK: - Lifecycle

    /// Begin playback-only (listening) operation. Called once at connect, after
    /// `client.setExternalPlayback(true)` and `AudioSessionManager.ensureSessionActive()`. Attaches
    /// the source node, wires the core's mixed-output sink into the ring, and starts the engine so
    /// remote audio plays immediately.
    func startListening(client: VoiceCatClient) {
        self.client = client
        guard !isConnected else { return }
        isConnected = true
        ring.reset()

        // Wire the core's mixed-output sink into the ring (C function pointer, no captures). Stays
        // registered for the whole connection; the ring is drained by the source-node render block.
        let ringPtr = Unmanaged.passUnretained(self.ring).toOpaque()
        client.setMixedOutputSink({ user, pcm, spc, ch, _ in
            guard let user, let pcm else { return }
            let ring = Unmanaged<PCMRing>.fromOpaque(user).takeUnretainedValue()
            ring.write(pcm, count: spc * Int(ch))
        }, user: ringPtr)

        rebuild()
    }

    /// Tear down the engine and unhook the core sink. Called on disconnect.
    func stop() {
        guard isConnected else { return }
        micActive = false
        stopMicTimer()
        isConnected = false
        client?.setMixedOutputSink(nil, user: nil)
        engine.inputNode.removeTap(onBus: 0)
        if engine.isRunning { engine.stop() }
        logger.info("audio engine stopped — ring written=\(self.ring.debugTotalWritten) read=\(self.ring.debugTotalRead) samples")
        try? engine.inputNode.setVoiceProcessingEnabled(false)
        if let src = sourceNode {
            engine.detach(src)
            sourceNode = nil
        }
        ring.reset()
        micRing.reset()
        client = nil
    }

    // MARK: - Mic transitions

    /// Engage the mic: install the input tap and (if the preset wants it) VPIO. Called when the
    /// user joins voice, after the MIC stream (external_feed) is started.
    func startMic(streamId: UInt32, channels: UInt32) {
        micStreamId = streamId
        captureChannels = channels
        micActive = true
        rebuild()
    }

    /// Disengage the mic: remove the tap and VPIO, keep playback running for remaining remote audio.
    func stopMic() {
        guard micActive else { return }
        micActive = false
        rebuild()
    }

    /// Update the capture channel count (mono↔stereo) for the active mic and rebuild.
    func setCaptureChannels(_ channels: UInt32) {
        captureChannels = channels
        if let client, micStreamId != 0 {
            client.setCaptureChannels(streamId: micStreamId, channels: channels)
        }
        if micActive { rebuild() }
    }

    /// Re-apply the engine graph against the current AVAudioSession config (preset / route change).
    /// Safe to call when only listening — it just rebuilds the playback graph against the new route.
    func reconfigure() {
        guard isConnected else { return }
        rebuild()
    }

    // MARK: - Graph (re)build

    /// The single place that (re)builds and starts the engine graph. Deterministic: stop → set
    /// VPIO → (re)install the mic tap → start. The caller is responsible for having applied the
    /// AVAudioSession config (category/mode/route) first (`IOSAudioRouter.applyConfiguration`).
    private func rebuild() {
        guard isConnected else { return }
        // Stop the feed pump before touching the tap / ring so the timer (on micQueue) can't race
        // the ring reset in installMicTap. It is restarted at the end with the current channel count.
        stopMicTimer()
        if engine.isRunning { engine.stop() }
        engine.inputNode.removeTap(onBus: 0)

        let useVPIO = micActive && IOSAudioRouter.shared.currentConfigUsesVoiceProcessing
        do {
            try engine.inputNode.setVoiceProcessingEnabled(useVPIO)
        } catch {
            logger.error("setVoiceProcessingEnabled(\(useVPIO)) failed: \(error.localizedDescription)")
        }
        if useVPIO {
            // AGC is the one VPIO sub-stage iOS exposes; AEC+NS are bundled into the master switch.
            engine.inputNode.isVoiceProcessingAGCEnabled = IOSAudioRouter.shared.agcEnabled
        }

        // The source node must bind to the selected voice-processing output unit.
        rebuildSourceNode()
        if micActive { installMicTap() }

        engine.prepare()
        do {
            try engine.start()
            let inFmt = engine.inputNode.outputFormat(forBus: 0)
            let outFmt = engine.outputNode.outputFormat(forBus: 0)
            let route = AVAudioSession.sharedInstance().currentRoute.outputs
                .map { "\($0.portName)[\($0.portType.rawValue)]" }.joined(separator: ", ")
            logger.info("""
                engine started — mic=\(self.micActive) vpio=\(useVPIO) captureCh=\(self.captureChannels) \
                inFormat=\(inFmt) outputNode=\(outFmt) outputRoute=[\(route)]
                """)
        } catch {
            // Route changes can leave AVAudioSession inactive; retry once after reactivation.
            logger.error("engine start failed: \(error.localizedDescription) — attempting one-shot recovery")
            do {
                try AudioSessionManager.shared.ensureSessionActive()
            } catch {
                logger.error("recovery — session re-activate failed: \(error.localizedDescription)")
            }
            IOSAudioRouter.shared.applyConfiguration()
            if AudioSessionManager.shared.isActive {
                IOSAudioRouter.shared.applyA2dpSpeakerFallback()
            }
            do {
                try engine.start()
                logger.info("engine start succeeded after one-shot recovery")
            } catch {
                logger.error("engine start failed after recovery: \(error.localizedDescription)")
                // Not fatal — a subsequent route-change or AVAudioEngine configuration-change
                // notification will trigger recoverAudio() and re-attempt the rebuild.
            }
        }

        // Start the feed pump last, with the current channel count, so it never carries a stale
        // (frozen) channel count across a mono↔stereo switch.
        if micActive { startMicTimer() }
    }

    /// Detach any previous source node and attach a fresh one pulling mixed PCM from the ring.
    /// Rebuilt on every graph rebuild so it always connects against the current output unit (the
    /// VPIO state can change the output between rebuilds). Its format is route-independent —
    /// `mainMixerNode` adapts 48 kHz stereo to whatever the output route is.
    private func rebuildSourceNode() {
        if let old = sourceNode {
            engine.detach(old)
            sourceNode = nil
        }
        let ring = self.ring
        let scratch = self.renderScratch
        let scratchFrames = self.renderScratchFrames
        let src = AVAudioSourceNode(format: outFormat) { _, _, frameCount, ablPtr in
            let frames = Int(frameCount)
            let abl = UnsafeMutableAudioBufferListPointer(ablPtr)
            let n = min(frames, scratchFrames)
            let got = ring.read(into: scratch, count: n * 2) / 2  // interleaved stereo → frames
            let scale: Float = 1.0 / 32768.0
            for ch in 0..<abl.count {
                guard let base = abl[ch].mData?.assumingMemoryBound(to: Float.self) else { continue }
                for i in 0..<frames {
                    base[i] = i < got ? Float(scratch[i * 2 + min(ch, 1)]) * scale : 0
                }
            }
            return noErr
        }
        sourceNode = src
        engine.attach(src)
        engine.connect(src, to: engine.mainMixerNode, format: outFormat)
    }

    /// Install the mic tap: convert the input node's native format to 48 kHz int16 (mono or
    /// stereo per `captureChannels`) and write it to the pacing ring. The 20 ms feed pump
    /// (`startMicTimer`) releases steady 960-sample frames to `feedPcm` — see the mic-feed comment
    /// above for why the tap must NOT call feedPcm directly (it bursts packets → receiver flutter).
    /// Rebuilds the converter each time because the input format depends on the VPIO state and the
    /// active route.
    private func installMicTap() {
        guard client != nil else { return }
        // Fresh ring on every (re)install — a rebuild must not feed stale pre-roll into the new tap.
        // Safe here: the feed pump was stopped at the top of rebuild(), so no consumer is running.
        micRing.reset()
        let ring = micRing  // captured by the closure as a `let` — no self capture (see mic-feed comment)
        let inFormat = engine.inputNode.outputFormat(forBus: 0)
        guard inFormat.sampleRate > 0 else {
            logger.error("input format unavailable (\(inFormat)) — mic will not transmit")
            return
        }
        let targetCh = max(1, min(2, captureChannels))
        guard let target = AVAudioFormat(commonFormat: .pcmFormatInt16, sampleRate: 48000,
                                         channels: AVAudioChannelCount(targetCh), interleaved: true),
              let converter = AVAudioConverter(from: inFormat, to: target) else {
            logger.error("mic converter unavailable (in=\(inFormat), ch=\(targetCh)) — mic will not transmit")
            return
        }

        let chInt = Int(targetCh)
        engine.inputNode.installTap(onBus: 0, bufferSize: 960, format: inFormat) { buffer, _ in
            // Convert this tap buffer to 48 kHz int16. Output capacity scaled for any upsample.
            let ratio = target.sampleRate / buffer.format.sampleRate
            let outCap = AVAudioFrameCount(Double(buffer.frameLength) * ratio + 16)
            guard let outBuf = AVAudioPCMBuffer(pcmFormat: target, frameCapacity: outCap) else { return }
            var fed = false
            let status = converter.convert(to: outBuf, error: nil) { _, outStatus in
                if fed { outStatus.pointee = .noDataNow; return nil }
                fed = true
                outStatus.pointee = .haveData
                return buffer
            }
            guard status != .error, outBuf.frameLength > 0,
                  let chData = outBuf.int16ChannelData else { return }
            // int16 interleaved → channelData[0] is the interleaved buffer. Write the converter's
            // variable-length output to the pacing ring; the 20 ms feed pump releases steady
            // 960-sample frames to feedPcm so packets leave the core at a steady 20 ms cadence.
            ring.write(chData[0], count: Int(outBuf.frameLength) * chInt)
        }
    }

    // MARK: - Mic feed pump (paces feedPcm at a steady 20 ms cadence)

    /// Start the 20 ms feed pump. After priming a small cushion (`pumpState.targetFrames`), it
    /// releases ONE 960-sample frame per tick from `micRing` to `feedPcm`, so the core (which sends
    /// synchronously per captured frame) emits packets at a steady 20 ms — the cadence its receivers
    /// expect. The cushion is essential: the receiver's playout deliberately keeps near-zero
    /// buffering (low latency), so it tolerates a steady stream but not bursts; the iOS tap delivers
    /// ~2 frames at once, and without the cushion the pump runs at ~0 depth and underruns on every
    /// tap/timer phase beat (crackle). Recreated on every `rebuild()` so `ch` always reflects the
    /// current `captureChannels` (mono↔stereo switches). Captures only locals + the reference-type
    /// ring/client/state (no `self`, which is @MainActor).
    private func startMicTimer() {
        stopMicTimer()
        guard let client else { return }
        let ring = micRing
        let scratch = micDrainScratch
        let state = pumpState
        let sid = micStreamId
        let ch = max(1, min(2, Int(captureChannels)))
        let frameSamples = Self.micFrameSamplesPerChannel
        let full = frameSamples * ch
        let chU32 = UInt32(ch)
        // The ring was just reset in installMicTap, so the cushion must be refilled before sending.
        state.primed = false
        let feed: () -> Void = {
            _ = ring.read(into: scratch, count: full)  // caller guarantees a full frame is present
            _ = client.feedPcm(streamId: sid, pcm: scratch,
                               samplesPerChannel: frameSamples, channels: chU32)
        }
        let t = DispatchSource.makeTimerSource(queue: micQueue)
        t.schedule(deadline: .now(), repeating: .milliseconds(20), leeway: .milliseconds(2))
        t.setEventHandler {
            let frames = ring.availableSamples / full  // whole frames currently buffered
            if !state.primed {
                if frames < state.targetFrames { return }  // still filling the cushion (into silence)
                state.primed = true
            } else if frames == 0 {
                // Re-prime with a larger cushion; consuming a partial frame would lose samples.
                if state.targetFrames < PumpState.maxTargetFrames { state.targetFrames += 1 }
                state.primed = false
                return
            }
            feed()  // one steady frame per tick (frames >= 1 here)
            // Catch-up: if the backlog grew past the cushion (pump descheduled, or producer ran
            // ahead via a burst), release one extra frame to drain it and keep latency bounded.
            if frames - 1 > state.targetFrames + 1 { feed() }
        }
        t.resume()
        micTimer = t
    }

    private func stopMicTimer() {
        micTimer?.cancel()
        micTimer = nil
    }
}