//! Audio Worklet backend implementation.
//!
//! Available on WebAssembly with the `audioworklet` feature. Requires atomics support.
//! See the `audioworklet-beep` example for setup instructions.

use std::{
    fmt,
    sync::{
        atomic::{AtomicU64, Ordering},
        Arc,
    },
    time::Duration,
};

use futures_channel::mpsc;
use futures_util::StreamExt as _;
use js_sys::wasm_bindgen;
use wasm_bindgen::prelude::*;

use crate::{
    host::frames_to_duration,
    traits::{DeviceTrait, HostTrait, StreamTrait},
    BufferSize, ChannelCount, Data, DeviceDescription, DeviceDescriptionBuilder, DeviceDirection,
    DeviceId, Error, ErrorKind, FrameCount, InputCallbackInfo, OutputCallbackInfo,
    OutputStreamTimestamp, Sample, SampleFormat, SampleRate, StreamConfig, StreamInstant,
    SupportedBufferSize, SupportedStreamConfig, SupportedStreamConfigRange,
};

mod dependent_module;
use crate::dependent_module;

/// Content is false if the iterator is empty.
pub struct Devices(bool);

#[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;

// `Stream`'s fields are all Send+Sync (a channel sender plus atomics).
// Any JS-backed resource (e.g. `web_sys::Window`, `Closure`) must live as a local
// inside the `spawn_local` task in `build_*_stream_raw`.
pub struct Stream {
    command_tx: mpsc::UnboundedSender<Command>,
    current_time_bits: Arc<AtomicU64>,
    buffer_size_frames: Arc<AtomicU64>,
}

crate::assert_stream_send!(Stream);
crate::assert_stream_sync!(Stream);

enum Command {
    Play,
    Pause,
}

/// How often the main thread re-reads `outputLatency` to publish it to the worklet.
const LATENCY_POLL_INTERVAL: Duration = Duration::from_millis(500);

pub use crate::iter::{SupportedInputConfigs, SupportedOutputConfigs};

// https://webaudio.github.io/web-audio-api/#dom-audioworkletnode-audioworkletnode
const MIN_CHANNELS: ChannelCount = 1;
const MAX_CHANNELS: ChannelCount = 32;

// https://webaudio.github.io/web-audio-api/#supported-sample-rates
const MIN_SAMPLE_RATE: SampleRate = 3_000;
const MAX_SAMPLE_RATE: SampleRate = 768_000;

// https://webaudio.github.io/web-audio-api/#audio-processing-model
const SUPPORTED_SAMPLE_FORMAT: SampleFormat = SampleFormat::F32;

// https://webaudio.github.io/web-audio-api/#render-quantum-size
const DEFAULT_RENDER_SIZE: u64 = 128;

fn render_quantum_size_supported() -> bool {
    (|| -> Option<bool> {
        let global = js_sys::global();
        let ctor = js_sys::Reflect::get(&global, &JsValue::from("AudioContext")).ok()?;
        let proto = js_sys::Reflect::get(&ctor, &JsValue::from("prototype")).ok()?;
        js_sys::Reflect::has(&proto, &JsValue::from("renderQuantumSize")).ok()
    })()
    .unwrap_or(false)
}

fn supported_render_quantum_range(sample_rate: SampleRate) -> SupportedBufferSize {
    // https://webaudio.github.io/web-audio-api/#supported-render-quantum-sizes
    if render_quantum_size_supported() {
        SupportedBufferSize::Range {
            min: 1,
            max: sample_rate.saturating_mul(6),
        }
    } else {
        SupportedBufferSize::Range {
            min: DEFAULT_RENDER_SIZE as FrameCount,
            max: DEFAULT_RENDER_SIZE as FrameCount,
        }
    }
}

impl Host {
    pub fn new() -> Result<Self, Error> {
        if Self::is_available() {
            Ok(Host)
        } else {
            Err(Error::with_message(
                ErrorKind::HostUnavailable,
                "AudioWorklet is not available",
            ))
        }
    }
}

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

    fn is_available() -> bool {
        if let Some(window) = web_sys::window() {
            let has_audio_worklet =
                js_sys::Reflect::has(&window, &JsValue::from_str("AudioWorklet")).unwrap_or(false);

            let cross_origin_isolated =
                js_sys::Reflect::get(&window, &JsValue::from_str("crossOriginIsolated"))
                    .ok()
                    .and_then(|v| v.as_bool())
                    .unwrap_or(false);

            has_audio_worklet && cross_origin_isolated
        } else {
            false
        }
    }

    fn devices(&self) -> Result<Self::Devices, Error> {
        Devices::new()
    }

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

    fn default_output_device(&self) -> Option<Self::Device> {
        if Self::is_available() {
            Some(Device)
        } else {
            None
        }
    }
}

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

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

    fn description(&self) -> Result<DeviceDescription, Error> {
        Ok(DeviceDescriptionBuilder::new("Default Device")
            .direction(DeviceDirection::Output)
            .build())
    }

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

    fn supported_input_configs(&self) -> Result<Self::SupportedInputConfigs, Error> {
        // TODO
        Ok(Vec::new().into_iter())
    }

    fn supported_output_configs(&self) -> Result<Self::SupportedOutputConfigs, Error> {
        // In actuality the number of supported channels cannot be fully known until
        // the browser attempts to initialized the AudioWorklet.

        let configs: Vec<_> = (MIN_CHANNELS..=MAX_CHANNELS)
            .flat_map(|channels| {
                crate::COMMON_SAMPLE_RATES
                    .iter()
                    .copied()
                    .filter(|&r| (MIN_SAMPLE_RATE..=MAX_SAMPLE_RATE).contains(&r))
                    .map(move |rate| SupportedStreamConfigRange {
                        channels,
                        min_sample_rate: rate,
                        max_sample_rate: rate,
                        buffer_size: supported_render_quantum_range(rate),
                        sample_format: SUPPORTED_SAMPLE_FORMAT,
                    })
            })
            .collect();
        Ok(configs.into_iter())
    }

    fn default_input_config(&self) -> Result<SupportedStreamConfig, Error> {
        Err(Error::with_message(
            ErrorKind::UnsupportedOperation,
            "Device does not support input",
        ))
    }

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

        Ok(config)
    }

    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,
    {
        Err(Error::with_message(
            ErrorKind::UnsupportedOperation,
            "Device does not support input",
        ))
    }

    /// Create an output stream.
    ///
    /// # Async completion
    ///
    /// This function returns `Ok` synchronously once the [`AudioContext`] is created, before the
    /// AudioWorklet module has been loaded or the [`AudioWorkletNode`] has been initialized. The
    /// actual worklet setup runs asynchronously via [`wasm_bindgen_futures::spawn_local`]. If
    /// setup fails (e.g. `add_module` or `AudioWorkletNode` construction throws), the error is
    /// delivered to `error_callback` after the caller already holds a [`Stream`]. There is no
    /// way to surface such errors synchronously given the Web Audio API's design.
    ///
    /// [`AudioContext`]: web_sys::AudioContext
    /// [`AudioWorkletNode`]: web_sys::AudioWorkletNode
    fn build_output_stream_raw<D, E>(
        &self,
        config: StreamConfig,
        sample_format: SampleFormat,
        data_callback: D,
        mut 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: the polled outputLatency can drop when the
        // page calls `setSinkId()` to switch output devices, pulling `playback` backward.
        let mut data_callback = crate::host::monotonic_output_callback(data_callback);
        if config.channels > MAX_CHANNELS {
            return Err(Error::with_message(
                ErrorKind::UnsupportedConfig,
                format!(
                    "Channel count {} exceeds the maximum of {MAX_CHANNELS}",
                    config.channels
                ),
            ));
        }
        if sample_format != SUPPORTED_SAMPLE_FORMAT {
            return Err(Error::with_message(
                ErrorKind::UnsupportedConfig,
                format!(
                    "Sample format {sample_format} is not supported; required format is {SUPPORTED_SAMPLE_FORMAT}"
                ),
            ));
        }
        if !(MIN_SAMPLE_RATE..=MAX_SAMPLE_RATE).contains(&config.sample_rate) {
            return Err(Error::with_message(
                ErrorKind::UnsupportedConfig,
                format!(
                    "Sample rate {} Hz is not in the supported range {MIN_SAMPLE_RATE}..={MAX_SAMPLE_RATE} Hz",
                    config.sample_rate
                ),
            ));
        }

        if let BufferSize::Fixed(n) = config.buffer_size {
            if let SupportedBufferSize::Range { min, max } =
                supported_render_quantum_range(config.sample_rate)
            {
                if !(min..=max).contains(&n) {
                    return Err(Error::with_message(
                        ErrorKind::UnsupportedConfig,
                        format!(
                            "Buffer size {n} is not in the supported render quantum range {min}..={max}"
                        ),
                    ));
                }
            }
        }

        let stream_opts = web_sys::AudioContextOptions::new();
        stream_opts.set_sample_rate(config.sample_rate as f32);
        if let BufferSize::Fixed(n) = config.buffer_size {
            let _ = js_sys::Reflect::set(
                stream_opts.as_ref(),
                &JsValue::from_str("renderSizeHint"),
                &JsValue::from_f64(n as f64),
            );
        }

        let audio_context =
            web_sys::AudioContext::new_with_context_options(&stream_opts).map_err(|_| {
                Error::with_message(
                    ErrorKind::UnsupportedConfig,
                    "Failed to create audio context",
                )
            })?;

        let destination = audio_context.destination();

        // Chrome rounds renderSizeHint to a power of two; read back the actual quantum.
        let actual_render_quantum =
            js_sys::Reflect::get(audio_context.as_ref(), &JsValue::from("renderQuantumSize"))
                .ok()
                .and_then(|v| v.as_f64())
                .map(|v| v as u64);

        if config.channels as u32 > destination.max_channel_count() {
            return Err(Error::with_message(
                ErrorKind::UnsupportedConfig,
                format!(
                    "Channel count {} exceeds the destination's maximum of {}",
                    config.channels,
                    destination.max_channel_count()
                ),
            ));
        }
        destination.set_channel_count(config.channels as u32);

        let initial_quantum = actual_render_quantum.unwrap_or(match config.buffer_size {
            BufferSize::Fixed(n) => n as u64,
            BufferSize::Default => DEFAULT_RENDER_SIZE,
        });
        let buffer_size_frames = Arc::new(AtomicU64::new(initial_quantum));
        let buffer_size_frames_cb = buffer_size_frames.clone();

        let current_time_bits = Arc::new(AtomicU64::new(audio_context.current_time().to_bits()));
        let current_time_bits_cb = current_time_bits.clone();
        let current_time_bits_init = current_time_bits.clone();

        let (command_tx, mut command_rx) = mpsc::unbounded::<Command>();

        // The worklet realm cannot read AudioContext properties, so share the value via an atomic.
        let latency_nanos = Arc::new(AtomicU64::new(total_latency_nanos(&audio_context)));
        let latency_nanos_cb = latency_nanos.clone();
        let ctx = audio_context.clone();
        wasm_bindgen_futures::spawn_local(async move {
            let result: Result<(), JsValue> = async move {
                let mod_url = dependent_module!("worklet.js")?;
                wasm_bindgen_futures::JsFuture::from(ctx.audio_worklet()?.add_module(&mod_url)?)
                    .await?;

                let options = web_sys::AudioWorkletNodeOptions::new();

                let js_array = js_sys::Array::new();
                js_array.push(&JsValue::from_f64(destination.channel_count() as _));

                options.set_output_channel_count(&js_array);
                options.set_number_of_inputs(0);

                options.set_processor_options(Some(&js_sys::Array::of3(
                    &wasm_bindgen::module(),
                    &wasm_bindgen::memory(),
                    &WasmAudioProcessor::new(Box::new(
                        move |interleaved_data, frame_size, sample_rate, now| {
                            buffer_size_frames_cb.store(frame_size as u64, Ordering::Relaxed);
                            current_time_bits_cb.store(now.to_bits(), Ordering::Relaxed);
                            let data = interleaved_data.as_mut_ptr() as *mut ();
                            let mut data = unsafe {
                                Data::from_parts(data, interleaved_data.len(), sample_format)
                            };

                            let callback = StreamInstant::from_secs_f64(now);
                            let buffer_duration =
                                frames_to_duration(frame_size as FrameCount, sample_rate);
                            let latency =
                                Duration::from_nanos(latency_nanos_cb.load(Ordering::Relaxed));
                            let playback = callback + (buffer_duration + latency);
                            let timestamp = OutputStreamTimestamp { callback, playback };
                            let info = OutputCallbackInfo { timestamp };
                            (data_callback)(&mut data, &info);
                        },
                    ))
                    .pack()
                    .into(),
                )));
                // This name 'CpalProcessor' must match the name registered in worklet.js
                let audio_worklet_node =
                    web_sys::AudioWorkletNode::new_with_options(&ctx, "CpalProcessor", &options)?;

                audio_worklet_node.connect_with_audio_node(&destination)?;
                Ok(())
            }
            .await;

            if let Err(err) = result {
                let message = err
                    .as_string()
                    .unwrap_or_else(|| "Failed to initialize audio worklet".to_string());
                error_callback(Error::with_message(
                    ErrorKind::UnsupportedOperation,
                    message,
                ));

                // Close AudioContext and exit; dropping command_rx closes the channel,
                // so subsequent play()/pause() calls return HostUnavailable.
                let _ = audio_context.close();
                return;
            }

            current_time_bits_init.store(audio_context.current_time().to_bits(), Ordering::Relaxed);

            // outputLatency can change at runtime (e.g. an output-device switch) but is only
            // readable on the main thread, so poll it here and publish it to the worklet via the
            // shared atomic.
            let _latency_poller = web_sys::window().and_then(|window| {
                let poll_ctx = audio_context.clone();
                let poll_latency = latency_nanos.clone();
                let closure = Closure::<dyn FnMut()>::new(move || {
                    poll_latency.store(total_latency_nanos(&poll_ctx), Ordering::Relaxed);
                });
                window
                    .set_interval_with_callback_and_timeout_and_arguments_0(
                        closure.as_ref().unchecked_ref(),
                        LATENCY_POLL_INTERVAL.as_millis() as i32,
                    )
                    .ok()
                    .map(|interval_id| LatencyPoller {
                        window,
                        interval_id,
                        _closure: closure,
                    })
            });

            // Process play/pause commands from any thread until Stream is dropped.
            // Dropping Stream closes command_tx, which terminates this loop.
            while let Some(cmd) = command_rx.next().await {
                match cmd {
                    Command::Play => {
                        if audio_context.resume().is_err() {
                            error_callback(Error::with_message(
                                ErrorKind::DeviceNotAvailable,
                                "Failed to resume audio context",
                            ));
                        }
                    }
                    Command::Pause => {
                        if audio_context.suspend().is_err() {
                            error_callback(Error::with_message(
                                ErrorKind::DeviceNotAvailable,
                                "Failed to suspend audio context",
                            ));
                        }
                    }
                }
            }

            // Stream dropped: close the AudioContext on the main thread.
            let _ = audio_context.close();
        });

        Ok(Self::Stream {
            command_tx,
            current_time_bits,
            buffer_size_frames,
        })
    }
}

impl StreamTrait for Stream {
    fn buffer_size(&self) -> Result<FrameCount, Error> {
        Ok(self.buffer_size_frames.load(Ordering::Relaxed) as FrameCount)
    }

    fn play(&self) -> Result<(), Error> {
        self.command_tx.unbounded_send(Command::Play).map_err(|_| {
            Error::with_message(
                ErrorKind::HostUnavailable,
                "audio worklet initialization failed",
            )
        })
    }

    fn pause(&self) -> Result<(), Error> {
        self.command_tx.unbounded_send(Command::Pause).map_err(|_| {
            Error::with_message(
                ErrorKind::HostUnavailable,
                "audio worklet initialization failed",
            )
        })
    }

    fn now(&self) -> StreamInstant {
        StreamInstant::from_secs_f64(f64::from_bits(
            self.current_time_bits.load(Ordering::Relaxed),
        ))
    }
}

impl Iterator for Devices {
    type Item = Device;

    fn next(&mut self) -> Option<Self::Item> {
        if self.0 {
            self.0 = false;
            Some(Device)
        } else {
            None
        }
    }
}

type AudioProcessorCallback = Box<dyn FnMut(&mut [f32], u32, u32, f64)>;

/// WasmAudioProcessor provides an interface for the JavaScript code
/// running in the AudioWorklet to interact with Rust.
#[wasm_bindgen]
pub struct WasmAudioProcessor {
    interleaved_buffer: Vec<f32>,
    // Passes in an interleaved scratch buffer, frame size, sample rate, and current time.
    callback: AudioProcessorCallback,
}

impl WasmAudioProcessor {
    pub fn new(callback: AudioProcessorCallback) -> Self {
        Self {
            interleaved_buffer: Vec::new(),
            callback,
        }
    }
}

#[wasm_bindgen]
impl WasmAudioProcessor {
    pub fn process(
        &mut self,
        channels: u32,
        frame_size: u32,
        sample_rate: u32,
        current_time: f64,
    ) -> u32 {
        let frame_size = frame_size as usize;

        // Ensure there's enough space in the output buffer
        // This likely only occurs once, or very few times.
        let interleaved_buffer_size = channels as usize * frame_size;
        self.interleaved_buffer.resize(
            interleaved_buffer_size.max(self.interleaved_buffer.len()),
            f32::EQUILIBRIUM,
        );

        self.interleaved_buffer[..interleaved_buffer_size].fill(f32::EQUILIBRIUM);

        (self.callback)(
            &mut self.interleaved_buffer[..interleaved_buffer_size],
            frame_size as u32,
            sample_rate,
            current_time,
        );

        // Returns a pointer to the raw interleaved buffer to Javascript so
        // it can deinterleave it into the output buffers.
        //
        // Deinterleaving is done on the Javascript side because it's simpler and it may be faster.
        // Doing it this way avoids an extra copy and the JS deinterleaving code
        // is likely heavily optimized by the browser's JS engine,
        // although I have not tested that assumption.
        self.interleaved_buffer.as_mut_ptr() as _
    }

    /// Converts this `WasmAudioProcessor` into a raw pointer (as `usize`) for FFI use.
    ///
    /// Transfers ownership of the processor to the caller. The returned pointer must be passed to
    /// [`unpack`] exactly once. Failing to call [`unpack`] will leak the allocation.
    ///
    /// [`unpack`]: Self::unpack
    pub fn pack(self) -> usize {
        Box::into_raw(Box::new(self)) as usize
    }
    /// # Safety
    ///
    /// The `val` parameter must be a value previously returned by `Self::pack`.
    /// It must not have already been unpacked or deallocated, and must not be used after this call.
    /// Using an invalid or already-consumed pointer will result in undefined behavior.
    pub unsafe fn unpack(val: usize) -> Self {
        *Box::from_raw(val as *mut _)
    }
}

/// Drives a `setInterval` that refreshes the shared output-latency value.
struct LatencyPoller {
    window: web_sys::Window,
    interval_id: i32,
    _closure: Closure<dyn FnMut()>,
}

impl Drop for LatencyPoller {
    fn drop(&mut self) {
        self.window.clear_interval_with_handle(self.interval_id);
    }
}

/// Reads the playback buffer depth from a context.
fn total_latency_nanos(ctx: &web_sys::AudioContext) -> u64 {
    let read = |key: &str| {
        js_sys::Reflect::get(ctx.as_ref(), &JsValue::from(key))
            .ok()
            .and_then(|v| v.as_f64())
            .unwrap_or(0.0)
    };
    // `baseLatency` is fixed for the context lifetime; `outputLatency` can change.
    let secs = read("baseLatency") + read("outputLatency");
    if secs.is_finite() && secs > 0.0 {
        (secs * 1_000_000_000.0).round() as u64
    } else {
        0
    }
}
