//! CoreAudio implementation for iOS using AVAudioSession and RemoteIO Audio Units.

use std::{
    fmt,
    ptr::NonNull,
    sync::{
        atomic::{AtomicUsize, Ordering},
        Arc, Mutex,
    },
    time::Duration,
};

use coreaudio::audio_unit::{
    render_callback::{self, data},
    AudioUnit, Element, Scope,
};
use objc2_audio_toolbox::{kAudioOutputUnitProperty_EnableIO, kAudioUnitProperty_StreamFormat};
use objc2_avf_audio::AVAudioSession;
use objc2_core_audio_types::AudioBuffer;

use self::enumerate::{
    default_input_device, default_output_device, Devices, SupportedInputConfigs,
    SupportedOutputConfigs,
};
use super::{asbd_from_config, host_time_to_stream_instant};
use crate::{
    host::{
        equilibrium::fill_equilibrium, frames_to_duration, latch::Latch, try_emit_error,
        ErrorCallbackArc,
    },
    traits::{DeviceTrait, HostTrait, StreamTrait},
    BufferSize, ChannelCount, Data, DeviceDescription, DeviceDescriptionBuilder, DeviceId, Error,
    ErrorKind, FrameCount, InputCallbackInfo, InputStreamTimestamp, OutputCallbackInfo,
    OutputStreamTimestamp, ResultExt, SampleFormat, SampleRate, StreamConfig, StreamInstant,
    SupportedBufferSize, SupportedStreamConfig, SupportedStreamConfigRange,
};

pub mod enumerate;
mod session_event_manager;
use session_event_manager::SessionEventManager;

// These days the default of iOS is now F32 and no longer I16
const SUPPORTED_SAMPLE_FORMAT: SampleFormat = SampleFormat::F32;

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct Device;

impl fmt::Display for Device {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let desc = self.description().map_err(|_| fmt::Error)?;
        f.write_str(desc.name())
    }
}

pub struct Host;

impl Host {
    pub fn new() -> Result<Self, Error> {
        Ok(Host)
    }
}

impl HostTrait for Host {
    type Devices = Devices;
    type Device = Device;

    fn is_available() -> bool {
        true
    }

    fn devices(&self) -> Result<Self::Devices, Error> {
        Ok(Devices(vec![Device].into_iter()))
    }

    fn default_input_device(&self) -> Option<Self::Device> {
        default_input_device()
    }

    fn default_output_device(&self) -> Option<Self::Device> {
        default_output_device()
    }
}

impl Device {
    fn description(&self) -> Result<DeviceDescription, Error> {
        // Query AVAudioSession to determine actual input/output availability
        // SAFETY: AVAudioSession::sharedInstance() returns the global audio session singleton
        let direction = unsafe {
            let audio_session = AVAudioSession::sharedInstance();
            let input_channels = Some(audio_session.inputNumberOfChannels() as ChannelCount);
            let output_channels = Some(audio_session.outputNumberOfChannels() as ChannelCount);

            crate::device_description::direction_from_counts(input_channels, output_channels)
        };

        Ok(DeviceDescriptionBuilder::new("Default Device")
            .direction(direction)
            .build())
    }

    fn id(&self) -> Result<DeviceId, Error> {
        Ok(DeviceId::new(crate::platform::HostId::CoreAudio, "default"))
    }

    fn supported_input_configs(&self) -> Result<SupportedInputConfigs, Error> {
        Ok(get_supported_stream_configs(true))
    }

    fn supported_output_configs(&self) -> Result<SupportedOutputConfigs, Error> {
        Ok(get_supported_stream_configs(false))
    }

    fn default_input_config(&self) -> Result<SupportedStreamConfig, Error> {
        // Get the primary (exact channel count) config from supported configs
        let range = get_supported_stream_configs(true).next().ok_or_else(|| {
            Error::with_message(
                ErrorKind::UnsupportedConfig,
                "No supported input configuration",
            )
        })?;
        Ok(range
            .try_with_standard_sample_rate()
            .unwrap_or_else(|| range.with_max_sample_rate()))
    }

    fn default_output_config(&self) -> Result<SupportedStreamConfig, Error> {
        let range = get_supported_stream_configs(false)
            .max_by(|a, b| a.cmp_default_heuristics(b))
            .ok_or_else(|| {
                Error::with_message(
                    ErrorKind::UnsupportedConfig,
                    "No supported output configuration",
                )
            })?;
        Ok(range
            .try_with_standard_sample_rate()
            .unwrap_or_else(|| range.with_max_sample_rate()))
    }
}

#[derive(Clone, Copy, Default, PartialEq, Eq)]
enum PlaybackState {
    #[default]
    Stopped,
    Playing,
    Interrupted,
}

struct StreamInner {
    state: PlaybackState,
    audio_unit: AudioUnit,
}

impl StreamInner {
    fn stop_for_interruption(&mut self) {
        if self.state != PlaybackState::Playing {
            return;
        }
        // Unlike pause(), the OS is the one that actually halted the unit here; this
        // call only resyncs AudioUnit's own bookkeeping, so its result carries nothing.
        let _ = self.audio_unit.stop();
        self.state = PlaybackState::Interrupted;
    }

    fn resume_after_interruption(&mut self) {
        if self.state != PlaybackState::Interrupted {
            return;
        }
        if self.audio_unit.start().is_ok() {
            self.state = PlaybackState::Playing;
        }
    }
}

impl DeviceTrait for Device {
    type SupportedInputConfigs = SupportedInputConfigs;
    type SupportedOutputConfigs = SupportedOutputConfigs;
    type Stream = Stream;

    fn description(&self) -> Result<DeviceDescription, Error> {
        Device::description(self)
    }

    fn id(&self) -> Result<DeviceId, Error> {
        Device::id(self)
    }

    fn supported_input_configs(&self) -> Result<Self::SupportedInputConfigs, Error> {
        Device::supported_input_configs(self)
    }

    fn supported_output_configs(&self) -> Result<Self::SupportedOutputConfigs, Error> {
        Device::supported_output_configs(self)
    }

    fn default_input_config(&self) -> Result<SupportedStreamConfig, Error> {
        Device::default_input_config(self)
    }

    fn default_output_config(&self) -> Result<SupportedStreamConfig, Error> {
        Device::default_output_config(self)
    }

    fn build_input_stream_raw<D, E>(
        &self,
        config: StreamConfig,
        sample_format: SampleFormat,
        data_callback: D,
        error_callback: E,
        _timeout: Option<Duration>,
    ) -> Result<Self::Stream, Error>
    where
        D: FnMut(&Data, &InputCallbackInfo) + Send + 'static,
        E: FnMut(Error) + Send + 'static,
    {
        crate::validate_stream_config(&config)?;
        // Keep `capture` monotonic: a route change (e.g. the user connects a headset) can
        // raise the input buffer depth, pulling `capture` backward.
        let data_callback = crate::host::monotonic_input_callback(data_callback);
        // Configure buffer size and create audio unit
        let mut audio_unit = setup_stream_audio_unit(config, sample_format, true)?;
        // Buffer depth used to offset capture timestamps. AVAudioSession auto-reroutes to a new device,
        // which changes this depth, and is then refreshed by the session event manager.
        let latency_frames = Arc::new(AtomicUsize::new(input_latency_frames()));

        let error_callback: ErrorCallbackArc = Arc::new(Mutex::new(error_callback));
        let session_error_callback = error_callback.clone();
        let session_latency_frames = latency_frames.clone();

        // Set up input callback
        setup_input_callback(
            &mut audio_unit,
            sample_format,
            config.sample_rate,
            latency_frames,
            data_callback,
            move |e| {
                let _ = try_emit_error(&error_callback, e);
            },
        )?;

        let inner = Arc::new(Mutex::new(StreamInner {
            state: PlaybackState::default(),
            audio_unit,
        }));
        let session_manager = SessionEventManager::new(
            session_error_callback,
            Latch::new(),
            Some((session_latency_frames, true)),
            Arc::downgrade(&inner),
        );
        let stream = Stream {
            inner,
            session_manager,
        };
        stream.signal_ready();
        Ok(stream)
    }

    /// Create an output stream.
    fn build_output_stream_raw<D, E>(
        &self,
        config: StreamConfig,
        sample_format: SampleFormat,
        data_callback: D,
        error_callback: E,
        _timeout: Option<Duration>,
    ) -> Result<Self::Stream, Error>
    where
        D: FnMut(&mut Data, &OutputCallbackInfo) + Send + 'static,
        E: FnMut(Error) + Send + 'static,
    {
        crate::validate_stream_config(&config)?;
        // Keep `playback` monotonic: a route change (e.g. the user disconnects a headset
        // and audio falls back to the built-in speaker) can lower the output buffer
        // depth, pulling `playback` backward.
        let data_callback = crate::host::monotonic_output_callback(data_callback);
        // Configure buffer size and create audio unit
        let mut audio_unit = setup_stream_audio_unit(config, sample_format, false)?;

        // Buffer depth used to offset playback timestamps. AVAudioSession auto-reroutes to a new device,
        // which changes this depth, and is then refreshed by the session event manager.
        let latency_frames = Arc::new(AtomicUsize::new(output_latency_frames()));

        let error_callback: ErrorCallbackArc = Arc::new(Mutex::new(error_callback));
        let session_error_callback = error_callback.clone();
        let session_latency_frames = latency_frames.clone();

        // Set up output callback
        setup_output_callback(
            &mut audio_unit,
            sample_format,
            config.sample_rate,
            latency_frames,
            data_callback,
            move |e| {
                let _ = try_emit_error(&error_callback, e);
            },
        )?;

        let inner = Arc::new(Mutex::new(StreamInner {
            state: PlaybackState::default(),
            audio_unit,
        }));
        let session_manager = SessionEventManager::new(
            session_error_callback,
            Latch::new(),
            Some((session_latency_frames, false)),
            Arc::downgrade(&inner),
        );
        let stream = Stream {
            inner,
            session_manager,
        };
        stream.signal_ready();
        Ok(stream)
    }
}

pub struct Stream {
    inner: Arc<Mutex<StreamInner>>,
    session_manager: SessionEventManager,
}

impl Stream {
    fn signal_ready(&self) {
        self.session_manager.signal_ready();
    }
}

impl Drop for Stream {
    fn drop(&mut self) {
        // Ensure the latch is released even if signal_ready() was never called (error path).
        self.session_manager.signal_ready();
    }
}

impl StreamTrait for Stream {
    fn play(&self) -> Result<(), Error> {
        let mut stream = self.inner.lock().map_err(|_| {
            Error::with_message(ErrorKind::StreamInvalidated, "Stream lock poisoned")
        })?;
        if stream.state != PlaybackState::Playing {
            stream
                .audio_unit
                .start()
                .context("Failed to start audio unit")?;
            stream.state = PlaybackState::Playing;
        }
        Ok(())
    }

    fn pause(&self) -> Result<(), Error> {
        let mut stream = self.inner.lock().map_err(|_| {
            Error::with_message(ErrorKind::StreamInvalidated, "Stream lock poisoned")
        })?;
        if stream.state == PlaybackState::Playing {
            stream
                .audio_unit
                .stop()
                .context("Failed to stop audio unit")?;
        }
        stream.state = PlaybackState::Stopped;
        Ok(())
    }

    fn now(&self) -> StreamInstant {
        let m_host_time = unsafe { mach2::mach_time::mach_absolute_time() };
        host_time_to_stream_instant(m_host_time).expect("mach_timebase_info failed")
    }

    fn buffer_size(&self) -> Result<FrameCount, Error> {
        Ok(get_device_buffer_frames() as FrameCount)
    }
}

fn create_audio_unit() -> Result<AudioUnit, coreaudio::Error> {
    AudioUnit::new_uninitialized(coreaudio::audio_unit::IOType::RemoteIO)
}

fn configure_for_recording(audio_unit: &mut AudioUnit) -> Result<(), coreaudio::Error> {
    // Enable mic recording
    let enable_input = 1u32;
    audio_unit.set_property(
        kAudioOutputUnitProperty_EnableIO,
        Scope::Input,
        Element::Input,
        Some(&enable_input),
    )?;

    // Disable output
    let disable_output = 0u32;
    audio_unit.set_property(
        kAudioOutputUnitProperty_EnableIO,
        Scope::Output,
        Element::Output,
        Some(&disable_output),
    )?;

    Ok(())
}

/// Configure AVAudioSession with the requested buffer size.
///
/// Note: iOS may not honor the exact request due to system constraints.
fn set_audio_session_buffer_size(
    buffer_size: FrameCount,
    sample_rate: SampleRate,
) -> Result<(), Error> {
    // SAFETY: AVAudioSession::sharedInstance() returns the global audio session singleton
    let audio_session = unsafe { AVAudioSession::sharedInstance() };

    // Calculate preferred buffer duration in seconds
    let buffer_duration = buffer_size as f64 / sample_rate as f64;

    // Set the preferred IO buffer duration
    // SAFETY: setPreferredIOBufferDuration_error is safe to call with valid duration
    unsafe {
        audio_session
            .setPreferredIOBufferDuration_error(buffer_duration)
            .map_err(|e| {
                Error::with_message(
                    ErrorKind::UnsupportedConfig,
                    format!("Failed to set preferred I/O buffer duration: {e}"),
                )
            })?;
    }

    Ok(())
}

/// Get the actual buffer size from AVAudioSession.
///
/// This queries the current IO buffer duration from AVAudioSession and converts
/// it to frames based on the current sample rate.
fn get_device_buffer_frames() -> usize {
    // SAFETY: AVAudioSession methods are safe to call on the singleton instance
    unsafe {
        let audio_session = AVAudioSession::sharedInstance();
        let buffer_duration = audio_session.IOBufferDuration();
        let sample_rate = audio_session.sampleRate();
        // Round: the duration comes from an integer frame count, so the product is a whole number
        // that floating-point can render as N.9999, which truncation would drop to N-1.
        (buffer_duration * sample_rate).round() as usize
    }
}

/// Total capture buffer depth in frames: the IO buffer plus the hardware input latency.
pub(super) fn input_latency_frames() -> usize {
    // SAFETY: AVAudioSession methods are safe to call on the singleton instance
    let extra = unsafe {
        let audio_session = AVAudioSession::sharedInstance();
        (audio_session.inputLatency() * audio_session.sampleRate()).round() as usize
    };
    get_device_buffer_frames() + extra
}

/// Total playback buffer depth in frames: the IO buffer plus the hardware output latency.
pub(super) fn output_latency_frames() -> usize {
    // SAFETY: AVAudioSession methods are safe to call on the singleton instance
    let extra = unsafe {
        let audio_session = AVAudioSession::sharedInstance();
        (audio_session.outputLatency() * audio_session.sampleRate()).round() as usize
    };
    get_device_buffer_frames() + extra
}

// Typical iOS hardware buffer frame limits according to Apple Technical Q&A QA1631.
const BUFFER_SIZE_MIN: FrameCount = 256;
const BUFFER_SIZE_MAX: FrameCount = 4096;

/// Get supported stream config ranges for input (is_input=true) or output (is_input=false).
fn get_supported_stream_configs(is_input: bool) -> std::vec::IntoIter<SupportedStreamConfigRange> {
    // SAFETY: AVAudioSession methods are safe to call on the singleton instance
    let (sample_rate, max_channels) = unsafe {
        let audio_session = AVAudioSession::sharedInstance();
        let sample_rate = audio_session.sampleRate() as SampleRate;
        let max_channels = if is_input {
            audio_session.inputNumberOfChannels() as ChannelCount
        } else {
            audio_session.outputNumberOfChannels() as ChannelCount
        };
        (sample_rate, max_channels)
    };

    let buffer_size = SupportedBufferSize::Range {
        min: BUFFER_SIZE_MIN,
        max: BUFFER_SIZE_MAX,
    };

    // For input, only return the exact channel count (no flexibility)
    // For output, support flexible channel counts up to the hardware maximum
    let min_channels = if is_input { max_channels } else { 1 };

    let configs: Vec<_> = (min_channels..=max_channels)
        .map(|channels| SupportedStreamConfigRange {
            channels,
            min_sample_rate: sample_rate,
            max_sample_rate: sample_rate,
            buffer_size,
            sample_format: SUPPORTED_SAMPLE_FORMAT,
        })
        .collect();

    configs.into_iter()
}

/// Setup audio unit with common configuration for input or output streams.
fn setup_stream_audio_unit(
    config: StreamConfig,
    sample_format: SampleFormat,
    is_input: bool,
) -> Result<AudioUnit, Error> {
    if let BufferSize::Fixed(buffer_size) = config.buffer_size {
        if !(BUFFER_SIZE_MIN..=BUFFER_SIZE_MAX).contains(&buffer_size) {
            return Err(Error::with_message(
                ErrorKind::UnsupportedConfig,
                format!(
                    "Buffer size {buffer_size} is not in the supported range \
                     {BUFFER_SIZE_MIN}..={BUFFER_SIZE_MAX}"
                ),
            ));
        }
        set_audio_session_buffer_size(buffer_size, config.sample_rate)?;
    }

    let mut audio_unit = create_audio_unit()?;

    if is_input {
        configure_for_recording(&mut audio_unit)?;
    }

    // Set the stream format in interleaved mode
    // For input: Output scope of Input element (data coming out of input)
    // For output: Input scope of Output element (data going into output)
    let (scope, element) = if is_input {
        (Scope::Output, Element::Input)
    } else {
        (Scope::Input, Element::Output)
    };

    let asbd = asbd_from_config(config, sample_format);
    audio_unit.set_property(kAudioUnitProperty_StreamFormat, scope, element, Some(&asbd))?;

    audio_unit.initialize()?;

    Ok(audio_unit)
}

/// Extract AudioBuffer and convert to Data, handling differences between input and output.
///
/// # Safety
///
/// Caller must ensure:
/// - `args.data.data` points to valid AudioBufferList
/// - For input: AudioBufferList has at least one buffer
/// - Buffer data remains valid for the callback duration
#[inline]
unsafe fn extract_audio_buffer(
    args: &render_callback::Args<data::Raw>,
    bytes_per_channel: usize,
    sample_format: SampleFormat,
    is_input: bool,
) -> (AudioBuffer, Data) {
    let buffer = if is_input {
        // Input: access through buffer array
        let first_buf_ptr = core::ptr::addr_of!((*args.data.data).mBuffers) as *const AudioBuffer;
        core::ptr::read_unaligned(first_buf_ptr)
    } else {
        // Output: direct access
        let buf_ptr = core::ptr::addr_of!((*args.data.data).mBuffers[0]);
        core::ptr::read_unaligned(buf_ptr)
    };

    let mut data_ptr = buffer.mData as *mut ();
    let mut len = buffer.mDataByteSize as usize / bytes_per_channel;

    // SAFETY: slice::from_raw_parts requires a non-null pointer.
    if data_ptr.is_null() {
        data_ptr = NonNull::dangling().as_ptr();
        len = 0;
    }

    let data = Data::from_parts(data_ptr, len, sample_format);

    (buffer, data)
}

/// Setup input callback with proper latency calculation.
fn setup_input_callback<D, E>(
    audio_unit: &mut AudioUnit,
    sample_format: SampleFormat,
    sample_rate: SampleRate,
    latency_frames: Arc<AtomicUsize>,
    mut data_callback: D,
    mut error_callback: E,
) -> Result<(), Error>
where
    D: FnMut(&Data, &InputCallbackInfo) + Send + 'static,
    E: FnMut(Error) + Send + 'static,
{
    let bytes_per_channel = sample_format.sample_size();
    type Args = render_callback::Args<data::Raw>;

    audio_unit.set_input_callback(move |args: Args| {
        // SAFETY: CoreAudio provides valid AudioBufferList for the callback duration
        let (buffer, data) =
            unsafe { extract_audio_buffer(&args, bytes_per_channel, sample_format, true) };

        let callback = match host_time_to_stream_instant(args.time_stamp.mHostTime) {
            Err(err) => {
                error_callback(err);
                return Err(());
            }
            Ok(cb) => cb,
        };

        // Refreshed on route changes by the session event manager; fall back to this buffer's
        // own frame count if the depth is unknown (zero).
        let latency_frames = match latency_frames.load(Ordering::Relaxed) {
            0 => {
                let channels = buffer.mNumberChannels as usize;
                data.len().checked_div(channels).unwrap_or(0)
            }
            n => n,
        };
        let delay = frames_to_duration(latency_frames as FrameCount, sample_rate);
        let capture = callback.checked_sub(delay).unwrap_or(StreamInstant::ZERO);
        let timestamp = InputStreamTimestamp { callback, capture };

        let info = InputCallbackInfo { timestamp };
        data_callback(&data, &info);
        Ok(())
    })?;

    Ok(())
}

/// Setup output callback with proper latency calculation.
fn setup_output_callback<D, E>(
    audio_unit: &mut AudioUnit,
    sample_format: SampleFormat,
    sample_rate: SampleRate,
    latency_frames: Arc<AtomicUsize>,
    mut data_callback: D,
    mut error_callback: E,
) -> Result<(), Error>
where
    D: FnMut(&mut Data, &OutputCallbackInfo) + Send + 'static,
    E: FnMut(Error) + Send + 'static,
{
    let bytes_per_channel = sample_format.sample_size();
    type Args = render_callback::Args<data::Raw>;

    audio_unit.set_render_callback(move |args: Args| {
        // SAFETY: CoreAudio provides valid AudioBufferList for the callback duration
        let (buffer, mut data) =
            unsafe { extract_audio_buffer(&args, bytes_per_channel, sample_format, false) };

        if !buffer.mData.is_null() {
            let bytes = unsafe {
                std::slice::from_raw_parts_mut(
                    buffer.mData as *mut u8,
                    buffer.mDataByteSize as usize,
                )
            };
            fill_equilibrium(bytes, sample_format);
        }

        let callback = match host_time_to_stream_instant(args.time_stamp.mHostTime) {
            Err(err) => {
                error_callback(err);
                return Err(());
            }
            Ok(cb) => cb,
        };

        // Refreshed on route changes by the session event manager; fall back to this buffer's
        // own frame count if the depth is unknown (zero).
        let latency_frames = match latency_frames.load(Ordering::Relaxed) {
            0 => {
                let channels = buffer.mNumberChannels as usize;
                data.len().checked_div(channels).unwrap_or(0)
            }
            n => n,
        };
        let delay = frames_to_duration(latency_frames as FrameCount, sample_rate);
        let playback = callback + delay;
        let timestamp = OutputStreamTimestamp { callback, playback };

        let info = OutputCallbackInfo { timestamp };
        data_callback(&mut data, &info);
        Ok(())
    })?;

    Ok(())
}

#[cfg(test)]
mod tests {
    use crate::{BufferSize, SampleRate, StreamConfig};

    #[test]
    fn test_ios_fixed_buffer_size() {
        let host = crate::default_host();
        let device = host.default_output_device().unwrap();

        let config = StreamConfig {
            channels: 2,
            sample_rate: SampleRate(48000),
            buffer_size: BufferSize::Fixed(512),
        };

        let result = device.build_output_stream(
            &config,
            |_data: &mut [f32], _info: &OutputCallbackInfo| {},
            |_err| {},
            None,
        );

        assert!(
            result.is_ok(),
            "BufferSize::Fixed should be supported on iOS via AVAudioSession"
        );
    }
}
