audiofox 1.1.0

// SPDX-License-Identifier: GPL-2.0-or-later WITH Classpath-exception-2.0
// Audiofox, the audio capturer!
// Copyright (c) 2026 foxberries

use alloc::boxed::Box;
use alloc::format;
use alloc::string::ToString;
use core::mem::{size_of, zeroed};
use core::ptr::{NonNull, addr_of_mut};
use core::slice;

use anyhow::Context;
use block2::StackBlock;
use dispatch2::{DispatchQueue, DispatchQueueAttr};
use log::{debug, info};
use objc2::__framework_prelude::Retained;
use objc2::AnyThread;
use objc2::rc::RetainedFromIterator;
use objc2_avf_audio::{AVAudioFormat, AVAudioPCMBuffer};
use objc2_core_audio::{
    AudioDeviceCreateIOProcIDWithBlock, AudioDeviceDestroyIOProcID, AudioDeviceIOBlock,
    AudioDeviceIOProcID, AudioDeviceStart, AudioDeviceStop, AudioHardwareCreateProcessTap,
    AudioHardwareDestroyAggregateDevice, AudioHardwareDestroyProcessTap, AudioObjectID,
    CATapDescription, kAudioAggregateDeviceIsPrivateKey, kAudioAggregateDeviceNameKey,
    kAudioAggregateDeviceTapListKey, kAudioAggregateDeviceUIDKey,
    kAudioDevicePropertyBufferFrameSize, kAudioDevicePropertyBufferFrameSizeRange,
    kAudioHardwarePropertyTranslatePIDToProcessObject, kAudioObjectSystemObject,
    kAudioSubTapUIDKey, kAudioTapPropertyFormat,
};
use objc2_core_audio_types::{
    AudioBufferList, AudioStreamBasicDescription, AudioTimeStamp, AudioValueRange,
};
use objc2_core_foundation::{CFDictionary, CFRetained};
use objc2_foundation::{NSArray, NSNumber, NSUUID};
use uuid::Uuid;

use crate::common::audiofox::AUDIOFOX;
use crate::common::callback::AudioCallback;
use crate::common::chunk::Chunk;
use crate::common::result::Result;
use crate::common::source::AudioSource;
use crate::common::stream::latency;
use crate::platform::mac::wrapper::{DictBuilder, create_aggregate, get_prop_data, set_prop_data};

/// Per-stream values used to initialize and maintain one or more links.
pub struct StreamState {
    tap_id: AudioObjectID,
    aggregate_device_id: AudioObjectID,
    device_proc_id: AudioDeviceIOProcID,
}

/// Initialize a new stream.
pub fn start_stream(
    uuid: Uuid,
    sources: &[AudioSource],
    callback: AudioCallback,
    frames: usize,
) -> Result<StreamState> {
    // Perform the core five init operations
    let audio_object_ids = sources_to_aoids(sources)?;
    let (tap_id, tap_uuid) = create_process_tap(uuid, &audio_object_ids)?;
    let aggregate_device_id = create_aggregate_device(&tap_uuid.to_string())?;
    let stream_props: AudioStreamBasicDescription = get_stream_props(tap_id)?;
    let buffer_size = set_buffer_size(aggregate_device_id, frames)?;

    // Wrap stream properties in AVAudioFormat for use in the process loop
    let audio_format = unsafe {
        AVAudioFormat::initWithStreamDescription(AVAudioFormat::alloc(), NonNull::from(&stream_props))
    }.context("Failed to create AVAudioFormat")?;

    // Define the IO block which will serve as the audio process loop for the duration of the
    // stream. Synchronous execution is guaranteed, so we use a StackBlock instead of an RcBlock
    let block = StackBlock::new(
        move |_in_now: NonNull<AudioTimeStamp>,
              in_input_data: NonNull<AudioBufferList>,
              _in_input_time: NonNull<AudioTimeStamp>,
              _out_output_data: NonNull<AudioBufferList>,
              _in_output_time: NonNull<AudioTimeStamp>| {
            // If we received a PCM buffer, process it
            if let Some(pcm_buffer) = unsafe {
                AVAudioPCMBuffer::initWithPCMFormat_bufferListNoCopy_deallocator(
                    AVAudioPCMBuffer::alloc(), &audio_format, in_input_data, None,
                )
            } {
                // Get the internal AudioBufferList
                let data = unsafe { pcm_buffer.audioBufferList().as_ref() };

                // Process each of the buffers contained in the list
                for i in 0..data.mNumberBuffers as usize {
                    let buffer = &data.mBuffers[i];

                    // Emit the Chunk
                    callback(Chunk::from(unsafe {
                        slice::from_raw_parts(
                            buffer.mData as *const f32,
                            buffer.mDataByteSize as usize / size_of::<f32>(),
                        )
                    }));
                }
            }
        },
    );

    // Create a new dispatch queue. We also request concurrent execution because IOBlock dispatch is
    // guaranteed to be synchronous by AudioDeviceCreateIOProcIDWithBlock. This results in slightly
    // decreased latency.
    let queue = DispatchQueue::new(AUDIOFOX, DispatchQueueAttr::concurrent());

    // Create the audio stream and get its ID
    let mut device_proc_id: AudioDeviceIOProcID = None;
    check_status!(unsafe {
        AudioDeviceCreateIOProcIDWithBlock(
            NonNull::from(&mut device_proc_id),
            aggregate_device_id,
            Some(&queue),
            AudioDeviceIOBlock::from(&raw const block as *mut _),
        )
    }).context("Failed to create audio process loop")?;

    // SAFETY: This call is safe because proc_id is implemented correctly
    check_status!(unsafe { AudioDeviceStart(aggregate_device_id, device_proc_id) })
        .context("Failed to start audio process loop")?;

    info!(
        "capturing at rate={}, channels={}, latency={:.2}ms",
        stream_props.mSampleRate, stream_props.mChannelsPerFrame, latency(buffer_size as usize)
    );

    Ok(StreamState { tap_id, aggregate_device_id, device_proc_id })
}

/// Close a stream and clean up its resources.
pub fn stop_stream(state: &StreamState) -> Result<()> {
    // SAFETY: Each of these calls are safe if LinkState is valid.
    if state.device_proc_id.is_some() {
        // Stop the audio stream
        check_status!(unsafe { AudioDeviceStop(state.aggregate_device_id, state.device_proc_id) })
            .context("Failed to stop the audio process loop")?;

        // Destroy the audio stream
        check_status!(unsafe { AudioDeviceDestroyIOProcID(state.aggregate_device_id, state.device_proc_id) })
            .context("Failed to destroy audio process loop")?;
    }
    // Destroy the aggregate device
    if state.aggregate_device_id != 0 {
        check_status!(unsafe { AudioHardwareDestroyAggregateDevice(state.aggregate_device_id) })
            .context("Failed to destroy aggregate device")?;
    }
    // Destroy the process tap
    if state.tap_id != 0 {
        check_status!(unsafe { AudioHardwareDestroyProcessTap(state.tap_id) })
            .context("Failed to destroy process tap")?;
    }
    Ok(())
}

/// Convert one or more [`AudioSource`] to one or more [`AudioObjectID`].
fn sources_to_aoids(sources: &[AudioSource]) -> Result<Box<[AudioObjectID]>> {
    sources.iter().map(|source| {
        let mut object_id: AudioObjectID = 0;
        get_prop_data(
            kAudioObjectSystemObject as AudioObjectID,
            &query!(kAudioHardwarePropertyTranslatePIDToProcessObject),
            qual!(source.id[0]),
            size_of::<AudioObjectID>() as u32,
            addr_of_mut!(object_id),
        ).context("Failed to translate PID to AudioObjectID")?;
        Ok(object_id)
    }).collect()
}

/// Create a process tap for one or more [`AudioObjectID`].
fn create_process_tap(
    uuid: Uuid,
    process_object_ids: &[AudioObjectID],
) -> Result<(AudioObjectID, Retained<NSUUID>)> {
    unsafe {
        // Create NSArray of NSNumbers containing the AudioObjectID
        let process_nums = process_object_ids.iter().map(|id| NSNumber::new_u32(*id));
        let process_array = NSArray::retained_from_iter(process_nums);

        // Initialize tap description with the process array
        let tap_desc = CATapDescription::initStereoMixdownOfProcesses(
            CATapDescription::alloc(),
            &process_array,
        );

        // Set the tap's UUID
        let tap_uuid = NSUUID::from_bytes(uuid.into_bytes());
        tap_desc.setUUID(&tap_uuid);

        // Create the process tap
        let mut tap_id: AudioObjectID = 0;
        check_status!(AudioHardwareCreateProcessTap(
            Some(&*tap_desc),
            &raw mut tap_id
        )).context("Failed to create process tap")?;

        Ok((tap_id, tap_uuid))
    }
}

/// Create an aggregate device connected to the process tap.
fn create_aggregate_device(tap_uid: &str) -> Result<AudioObjectID> {
    // Create the tap dictionary
    let tap_dict: CFRetained<CFDictionary> = DictBuilder::<1>::new()
        .insert(kAudioSubTapUIDKey, tap_uid)
        .build()
        .context("Failed to create tap dictionary")?;

    // Create the aggregate device dictionary
    let aggregate_dict = DictBuilder::<4>::new()
        // Set the name of the aggregate device.
        .insert(kAudioAggregateDeviceNameKey, AUDIOFOX)
        // Set the aggregate device UID to uniquely identify this device.
        .insert(kAudioAggregateDeviceUIDKey, &*format!("{AUDIOFOX}-{tap_uid}"))
        // Set the tap list, which are the processes we will tap for their audio.
        .insert(kAudioAggregateDeviceTapListKey, tap_dict)
        // Set the private flag, so Audiofox does not appear in places like Audio MIDI Setup.
        .insert(kAudioAggregateDeviceIsPrivateKey, true)
        .build()
        .context("Failed to create aggregate dictionary")?;

    // Create the aggregate device using the dictionary and return its AudioObjectID
    create_aggregate(&aggregate_dict).context("Failed to create aggregate device")
}

/// Get the audio format of a stream from its process tap.
fn get_stream_props(tap_id: AudioObjectID) -> Result<AudioStreamBasicDescription> {
    // Initializing all fields with zeroes is safe and in-line with Apple's documentation.
    let mut stream_props: AudioStreamBasicDescription = unsafe { zeroed() };

    // Get the audio stream props
    get_prop_data(
        tap_id,
        &query!(kAudioTapPropertyFormat),
        None,
        size_of::<AudioStreamBasicDescription>() as u32,
        addr_of_mut!(stream_props),
    ).context("Failed to get stream props")?;

    Ok(stream_props)
}

/// Set the buffer size of a stream to a user-defined value.
fn set_buffer_size(aggregate_device_id: AudioObjectID, frames: usize) -> Result<u32> {
    // Query the system for the valid buffer size range
    let mut buffer_range: AudioValueRange = unsafe { zeroed() }; // Safe, see above
    get_prop_data(
        aggregate_device_id,
        &query!(kAudioDevicePropertyBufferFrameSizeRange),
        None,
        size_of::<AudioValueRange>() as u32,
        addr_of_mut!(buffer_range),
    ).context("Failed to get buffer size range")?;

    // Get the minimum and maximum size
    let min_frames = buffer_range.mMinimum as u32;
    let max_frames = buffer_range.mMaximum as u32;

    debug!("valid buffer range: {} - {} frames", min_frames, max_frames);

    // Clamp the desired size to the bounds
    let desired_size = frames as u32;
    let buffer_size: u32 = desired_size.max(min_frames).min(max_frames);

    // Set the buffer size
    set_prop_data(
        aggregate_device_id,
        &query!(kAudioDevicePropertyBufferFrameSize),
        size_of::<u32>() as u32,
        &raw const buffer_size,
    ).context("Failed to set buffer size")?;

    Ok(buffer_size)
}