// Native system-audio capture for the custom screen-share picker. Without
// this path the picker has to fall back to getDisplayMedia whenever the
// user ticks "Mit System-Sound", because Chromium only wires audio into
// desktop captures that the OS picker produced. Here we grab the default
// render endpoint's loopback stream via WASAPI, convert it to 48kHz f32
// stereo, and ship the samples to the JS side through a Tauri Channel.
// An AudioWorklet on the frontend feeds them into a MediaStreamDestination
// so LiveKit publishes a plain ScreenShareAudio track.
//
// Windows-only for v1. macOS + Linux stubs return a clear error so the
// frontend can fall back cleanly on those platforms until their native
// paths ship (ScreenCaptureKit-audio / PipeWire).

#![allow(clippy::needless_return)]

use base64::Engine;
use serde::Serialize;
use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
use std::sync::{Arc, Mutex, OnceLock};
use std::thread;
use tauri::ipc::Channel;

// Output format we always deliver to the frontend. Picking a single fixed
// format means the AudioWorklet never has to renegotiate — it just assumes
// interleaved f32 stereo at 48kHz. WASAPI mix format is usually already
// this on Windows 10+, so the resample branch is rarely hit.
const OUTPUT_SAMPLE_RATE: u32 = 48_000;
const OUTPUT_CHANNELS: u16 = 2;

static NEXT_ID: AtomicU32 = AtomicU32::new(1);
static SESSIONS: OnceLock<Mutex<HashMap<u32, Session>>> = OnceLock::new();

fn sessions() -> &'static Mutex<HashMap<u32, Session>> {
    SESSIONS.get_or_init(|| Mutex::new(HashMap::new()))
}

struct Session {
    stop: Arc<AtomicBool>,
    handle: Option<thread::JoinHandle<()>>,
}

#[derive(Serialize, Clone)]
#[serde(rename_all = "camelCase")]
pub struct AudioFramePayload {
    pub capture_id: u32,
    pub sample_rate: u32,
    pub channels: u16,
    /// Interleaved little-endian f32 stereo samples, base64-encoded.
    /// Frontend decodes via `atob` → `Uint8Array` → `Float32Array` view.
    /// Base64 is used instead of a raw `Vec<f32>` because Tauri Channel
    /// serialises via JSON — a JSON array of floats balloons to ~2–3×
    /// the byte count, and at 48kHz stereo that's enough IPC traffic
    /// to matter.
    pub samples_base64: String,
}

/// Start a loopback capture of the default render endpoint and begin
/// streaming audio frames on the provided channel. Returns a numeric
/// capture id that must be handed to `stop_system_audio_capture` when
/// the share ends.
#[tauri::command]
pub fn start_system_audio_capture(
    channel: Channel<AudioFramePayload>,
) -> Result<u32, String> {
    #[cfg(target_os = "windows")]
    {
        let capture_id = NEXT_ID.fetch_add(1, Ordering::Relaxed);
        let stop = Arc::new(AtomicBool::new(false));
        let stop_clone = Arc::clone(&stop);

        let handle = thread::Builder::new()
            .name(format!("screen-audio-{capture_id}"))
            .spawn(move || {
                if let Err(err) =
                    windows_loopback::capture_loop(capture_id, channel, stop_clone)
                {
                    eprintln!("screen-audio {capture_id}: {err}");
                }
            })
            .map_err(|e| format!("failed to spawn audio thread: {e}"))?;

        sessions().lock().unwrap().insert(
            capture_id,
            Session {
                stop,
                handle: Some(handle),
            },
        );
        Ok(capture_id)
    }

    #[cfg(not(target_os = "windows"))]
    {
        // Keep the `channel` binding alive so Tauri doesn't complain about
        // an unused parameter on the non-Windows build.
        let _ = channel;
        Err("system audio capture only supported on Windows".into())
    }
}

/// Tear down the capture for the given id. Safe to call on a missing id
/// (no-op) so the JS side doesn't have to track whether the stop has
/// already been issued by the screen-share teardown path.
#[tauri::command]
pub fn stop_system_audio_capture(capture_id: u32) -> Result<(), String> {
    let session = sessions().lock().unwrap().remove(&capture_id);
    let Some(mut session) = session else {
        return Ok(());
    };
    session.stop.store(true, Ordering::Relaxed);
    if let Some(handle) = session.handle.take() {
        // Best-effort join — the capture loop polls `stop` every event
        // cycle (≤100ms) so this usually returns promptly. If the WASAPI
        // call is wedged we'd rather drop the handle than hang the stop.
        let _ = handle.join();
    }
    Ok(())
}

// ---------------------------------------------------------------------------
// Windows loopback implementation
// ---------------------------------------------------------------------------

#[cfg(target_os = "windows")]
mod windows_loopback {
    use super::*;
    use wasapi::{initialize_mta, Direction, SampleType, ShareMode};

    pub fn capture_loop(
        capture_id: u32,
        channel: Channel<AudioFramePayload>,
        stop: Arc<AtomicBool>,
    ) -> Result<(), String> {
        // COM must be initialised on every thread that touches WASAPI.
        // MTA is the right model for a background capture thread — STA
        // would require message pumping we don't want to add.
        initialize_mta()
            .ok()
            .map_err(|e| format!("initialize_mta: {e:?}"))?;

        let device = wasapi::get_default_device(&Direction::Render)
            .map_err(|e| format!("get_default_device: {e:?}"))?;
        let mut audio_client = device
            .get_iaudioclient()
            .map_err(|e| format!("get_iaudioclient: {e:?}"))?;

        // Use the mix format that Windows is already pushing to the
        // endpoint. Loopback capture won't convert for us — asking for a
        // fixed format here makes Initialize() fail on non-matching
        // hardware. We resample + channel-mix ourselves downstream.
        let mix_format = audio_client
            .get_mixformat()
            .map_err(|e| format!("get_mixformat: {e:?}"))?;
        let input_rate = mix_format.get_samplespersec();
        let input_channels = mix_format.get_nchannels();
        let bits_per_sample = mix_format.get_bitspersample();
        let block_align = mix_format.get_blockalign();
        let sample_type = mix_format.get_subformat().unwrap_or(SampleType::Int);

        let (def_time, _min_time) = audio_client
            .get_periods()
            .map_err(|e| format!("get_periods: {e:?}"))?;

        // Direction::Capture + loopback: WASAPI streams what Windows is
        // sending to the speakers instead of what an input device is
        // producing. Shared mode so we coexist with other apps.
        audio_client
            .initialize_client(
                &mix_format,
                def_time,
                &Direction::Capture,
                &ShareMode::Shared,
                true,
            )
            .map_err(|e| format!("initialize_client: {e:?}"))?;

        let h_event = audio_client
            .set_get_eventhandle()
            .map_err(|e| format!("set_get_eventhandle: {e:?}"))?;
        let capture_client = audio_client
            .get_audiocaptureclient()
            .map_err(|e| format!("get_audiocaptureclient: {e:?}"))?;

        audio_client
            .start_stream()
            .map_err(|e| format!("start_stream: {e:?}"))?;

        // Resampler state — last stereo frame from the previous buffer so
        // linear interpolation at the buffer boundary doesn't click.
        // Initialised to silence.
        let mut last_stereo: [f32; 2] = [0.0, 0.0];

        while !stop.load(Ordering::Relaxed) {
            // 100ms timeout lets the loop check the stop flag even when
            // the endpoint is silent (WASAPI doesn't signal the event at
            // all for pure-silence streams on some driver versions).
            if h_event.wait_for_event(100).is_err() {
                continue;
            }

            // Drain all packets available since the last wake — there
            // can be several queued if we were preempted.
            loop {
                if stop.load(Ordering::Relaxed) {
                    break;
                }
                let frames_available = match capture_client.get_next_nbr_frames() {
                    Ok(Some(n)) if n > 0 => n,
                    Ok(_) => break,
                    Err(e) => {
                        eprintln!(
                            "screen-audio {capture_id}: get_next_nbr_frames: {e:?}"
                        );
                        break;
                    }
                };

                let bytes_needed =
                    frames_available as usize * block_align as usize;
                let mut raw = vec![0u8; bytes_needed];
                if let Err(e) = capture_client.read_from_device(&mut raw) {
                    eprintln!(
                        "screen-audio {capture_id}: read_from_device: {e:?}"
                    );
                    break;
                }

                // Decode PCM into interleaved f32 at the device's native
                // rate + channel count.
                let decoded = decode_pcm(
                    &raw,
                    input_channels,
                    bits_per_sample,
                    &sample_type,
                );

                // Channel-fold → 2ch, then resample → 48kHz.
                let stereo = to_stereo(&decoded, input_channels);
                let out_samples = resample_linear_stereo(
                    &stereo,
                    input_rate,
                    OUTPUT_SAMPLE_RATE,
                    &mut last_stereo,
                );

                if out_samples.is_empty() {
                    continue;
                }

                // Pack f32s as little-endian bytes then base64. IPC-wise
                // this is ~1.3× the raw byte count versus 5–10× for a
                // JSON array of floats, which is the difference between
                // "fine" and "wastes a CPU core" at 48kHz stereo.
                let mut bytes = Vec::with_capacity(out_samples.len() * 4);
                for s in &out_samples {
                    bytes.extend_from_slice(&s.to_le_bytes());
                }
                let samples_b64 =
                    base64::engine::general_purpose::STANDARD.encode(&bytes);

                if channel
                    .send(AudioFramePayload {
                        capture_id,
                        sample_rate: OUTPUT_SAMPLE_RATE,
                        channels: OUTPUT_CHANNELS,
                        samples_base64: samples_b64,
                    })
                    .is_err()
                {
                    // Frontend went away — stop cleanly.
                    stop.store(true, Ordering::Relaxed);
                    break;
                }
            }
        }

        let _ = audio_client.stop_stream();
        Ok(())
    }

    // Convert a raw WASAPI buffer into interleaved f32 at the device's
    // native channel count. Handles the three formats that actually show
    // up on Windows render endpoints: f32 (most modern hardware), i16
    // (older onboard codecs), and i32 (pro audio interfaces). Anything
    // else falls through to zeros so a weird format doesn't crash the
    // share — the user will notice silence and can retry.
    fn decode_pcm(
        raw: &[u8],
        channels: u16,
        bits_per_sample: u16,
        sample_type: &SampleType,
    ) -> Vec<f32> {
        match (sample_type, bits_per_sample) {
            (SampleType::Float, 32) => {
                let mut out = Vec::with_capacity(raw.len() / 4);
                for chunk in raw.chunks_exact(4) {
                    out.push(f32::from_le_bytes([
                        chunk[0], chunk[1], chunk[2], chunk[3],
                    ]));
                }
                out
            }
            (SampleType::Int, 16) => {
                let scale = 1.0_f32 / (i16::MAX as f32);
                let mut out = Vec::with_capacity(raw.len() / 2);
                for chunk in raw.chunks_exact(2) {
                    let s = i16::from_le_bytes([chunk[0], chunk[1]]);
                    out.push(s as f32 * scale);
                }
                out
            }
            (SampleType::Int, 32) => {
                let scale = 1.0_f32 / (i32::MAX as f32);
                let mut out = Vec::with_capacity(raw.len() / 4);
                for chunk in raw.chunks_exact(4) {
                    let s = i32::from_le_bytes([
                        chunk[0], chunk[1], chunk[2], chunk[3],
                    ]);
                    out.push(s as f32 * scale);
                }
                out
            }
            _ => {
                // Unknown format — emit silence of the right frame count
                // so downstream math stays correct.
                let bytes_per_frame =
                    (bits_per_sample as usize / 8) * channels as usize;
                let frames = if bytes_per_frame == 0 {
                    0
                } else {
                    raw.len() / bytes_per_frame
                };
                vec![0.0; frames * channels as usize]
            }
        }
    }

    // Down- or up-mix to stereo. Surround layouts fold L+R only (center
    // + surrounds get dropped) which is the simplest defensible choice
    // for screen-share audio — most content is LR-centric and a proper
    // ITU-R BS.775 downmix would pull in matrix coefficients we'd rather
    // avoid in v1.
    fn to_stereo(interleaved: &[f32], channels: u16) -> Vec<f32> {
        if channels == 0 || interleaved.is_empty() {
            return Vec::new();
        }
        if channels == 2 {
            return interleaved.to_vec();
        }
        let ch = channels as usize;
        let frames = interleaved.len() / ch;
        let mut out = Vec::with_capacity(frames * 2);
        if channels == 1 {
            for i in 0..frames {
                let s = interleaved[i];
                out.push(s);
                out.push(s);
            }
        } else {
            for i in 0..frames {
                let base = i * ch;
                out.push(interleaved[base]);
                out.push(interleaved[base + 1]);
            }
        }
        out
    }

    // Linear-interpolation resampler for interleaved stereo f32. Not the
    // prettiest option theoretically, but at 44.1→48 the audible
    // artefacts stay below threshold for speech + game/music content. The
    // `last_stereo` state preserves the final frame across invocations so
    // the interpolation at the buffer boundary doesn't produce a click.
    fn resample_linear_stereo(
        input_stereo: &[f32],
        input_rate: u32,
        output_rate: u32,
        last_stereo: &mut [f32; 2],
    ) -> Vec<f32> {
        if input_stereo.is_empty() {
            return Vec::new();
        }
        if input_rate == output_rate {
            last_stereo[0] = input_stereo[input_stereo.len() - 2];
            last_stereo[1] = input_stereo[input_stereo.len() - 1];
            return input_stereo.to_vec();
        }
        let ratio = output_rate as f64 / input_rate as f64;
        let in_frames = input_stereo.len() / 2;
        let out_frames = (in_frames as f64 * ratio).floor() as usize;
        if out_frames == 0 {
            last_stereo[0] = input_stereo[input_stereo.len() - 2];
            last_stereo[1] = input_stereo[input_stereo.len() - 1];
            return Vec::new();
        }
        let mut out = Vec::with_capacity(out_frames * 2);

        let prev_l = last_stereo[0];
        let prev_r = last_stereo[1];

        for i in 0..out_frames {
            let src_pos = i as f64 / ratio;
            let src_frame = src_pos.floor() as i64;
            let frac = (src_pos - src_frame as f64) as f32;

            // `src_frame == -1` comes up for the very first output frame
            // when ratio > 1 — interpolate against the previous buffer's
            // final sample to bridge the two.
            let (l0, r0) = if src_frame < 0 {
                (prev_l, prev_r)
            } else {
                let idx = (src_frame as usize).min(in_frames - 1) * 2;
                (input_stereo[idx], input_stereo[idx + 1])
            };
            let next = ((src_frame + 1) as usize).min(in_frames - 1);
            let l1 = input_stereo[next * 2];
            let r1 = input_stereo[next * 2 + 1];

            out.push(l0 + (l1 - l0) * frac);
            out.push(r0 + (r1 - r0) * frac);
        }

        last_stereo[0] = input_stereo[input_stereo.len() - 2];
        last_stereo[1] = input_stereo[input_stereo.len() - 1];
        out
    }
}