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
    }

    /// 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
///     calls `client.feedPcm(micStreamId)`.
///
/// 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

    // 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)
    }

    // 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
        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()
        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 }
        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
        }

        // (Re)build the playback source node AFTER the VPIO state is set, so it connects against the
        // correct (voice-processed or plain) output unit — mirrors the proven original ordering.
        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 {
            logger.error("engine start failed: \(error.localizedDescription)")
        }
    }

    /// 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 feed it to the core. Rebuilds the converter each time
    /// because the input format depends on the VPIO state and the active route.
    private func installMicTap() {
        guard let client else { return }
        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 sid = micStreamId
        let c = client
        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.
            c.feedPcm(streamId: sid, pcm: chData[0],
                      samplesPerChannel: Int(outBuf.frameLength), channels: targetCh)
        }
    }
}