audiofox 1.0.2

// 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::cell::RefCell;
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::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,
    kAudioHardwarePropertyTranslatePIDToProcessObject, kAudioObjectSystemObject,
    kAudioSubTapUIDKey, kAudioTapPropertyFormat,
};
use objc2_core_audio_types::{AudioBufferList, AudioStreamBasicDescription, AudioTimeStamp};
use objc2_core_foundation::{CFDictionary, CFRetained};
use objc2_foundation::{NSArray, NSNumber, NSUUID};
use uuid::Uuid;

use crate::common::audiofox::AUDIOFOX;
use crate::common::builder::ChunkBuilder;
use crate::common::callback::AudioCallback;
use crate::common::chunk::Chunk;
use crate::common::result::Result;
use crate::common::source::AudioSource;
use crate::platform::mac::wrapper::{DictBuilder, create_aggregate, get_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,
) -> Result<StreamState> {
    // Perform the core four 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 audio_format = get_audio_format(tap_id)?;

    // Initialize the current chunk wrapped in a RefCell
    let builder = RefCell::from(ChunkBuilder::new_u8());

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

                // Get the ChunkBuilder
                let mut builder = builder.borrow_mut();

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

                    // SAFETY: This call is safe because the data is guaranteed to be well-aligned for u8
                    let sample_slice = unsafe {
                        slice::from_raw_parts(
                            buffer.mData as *const u8,
                            buffer.mDataByteSize as usize,
                        )
                    };

                    // Aggregate samples until we have enough to send a full chunk
                    if let Some(data) = builder.extend_from_slice(sample_slice) {
                        // If data is returned, we're at capacity
                        callback(Chunk::from(&*builder));
                        // Reset the ChunkBuilder with the overflow
                        builder.reset_with_overflow(data);
                    }
                }
            }
        },
    );

    // 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!("stream successfully started");

    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_audio_format(tap_id: AudioObjectID) -> Result<Retained<AVAudioFormat>> {
    // Initializing all fields with zeroes is safe and in-line with Apple's documentation.
    let mut stream_format: AudioStreamBasicDescription = unsafe { zeroed() };

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

    // stream_format is initialized at this point
    info!("capturing at rate:{}, channels:{}", stream_format.mSampleRate, stream_format.mChannelsPerFrame);

    // Wrap in AVAudioFormat for immediate use by start_streaming()
    let av_format = unsafe {
        AVAudioFormat::initWithStreamDescription(
            AVAudioFormat::alloc(),
            NonNull::from(&stream_format),
        )
    }.context("Failed to create AVAudioFormat")?;

    Ok(av_format)
}