//! WASAPI loopback capture for Windows.
//!
//! This module captures system audio (what's playing through speakers)
//! using Windows Audio Session API (WASAPI) in loopback mode.

use parking_lot::Mutex;
use std::sync::Arc;
use std::thread::{self, JoinHandle};
use tracing::{error, info, warn};
use windows::Win32::Media::Audio;
use windows::Win32::System::Com;

/// WASAPI loopback capture state.
pub struct WasapiLoopback {
    is_capturing: Arc<Mutex<bool>>,
    should_stop: Arc<Mutex<bool>>,
    capture_thread: Option<JoinHandle<()>>,
}

impl WasapiLoopback {
    /// Create a new WASAPI loopback capture instance.
    pub fn new() -> Self {
        Self {
            is_capturing: Arc::new(Mutex::new(false)),
            should_stop: Arc::new(Mutex::new(false)),
            capture_thread: None,
        }
    }

    /// Start capturing system audio in loopback mode.
    pub fn start_capture(&mut self, output_buffer: Arc<Mutex<Vec<f32>>>) -> Result<(), String> {
        if *self.is_capturing.lock() {
            warn!("WASAPI loopback already capturing");
            return Ok(());
        }

        *self.should_stop.lock() = false;
        *self.is_capturing.lock() = true;

        let is_capturing = Arc::clone(&self.is_capturing);
        let should_stop = Arc::clone(&self.should_stop);

        let handle = thread::spawn(move || {
            if let Err(e) = capture_loopback_audio(output_buffer, should_stop.clone()) {
                error!("WASAPI loopback capture error: {}", e);
            }
            *is_capturing.lock() = false;
            info!("WASAPI loopback capture stopped");
        });

        self.capture_thread = Some(handle);
        Ok(())
    }

    /// Stop capturing.
    pub fn stop_capture(&mut self) {
        *self.should_stop.lock() = true;
        if let Some(handle) = self.capture_thread.take() {
            let _ = handle.join();
        }
    }

    /// Check if currently capturing.
    pub fn is_capturing(&self) -> bool {
        *self.is_capturing.lock()
    }
}

impl Drop for WasapiLoopback {
    fn drop(&mut self) {
        self.stop_capture();
    }
}

/// Target sample rate for Whisper (16kHz)
const TARGET_SAMPLE_RATE: u32 = 16000;

/// Capture loopback audio from the default render device.
fn capture_loopback_audio(
    buffer: Arc<Mutex<Vec<f32>>>,
    should_stop: Arc<Mutex<bool>>,
) -> Result<(), String> {
    unsafe {
        // Initialize COM
        Com::CoInitializeEx(None, Com::COINIT_MULTITHREADED)
            .ok()
            .map_err(|e| format!("Failed to initialize COM: {}", e))?;

        // Create device enumerator
        let enumerator: Audio::IMMDeviceEnumerator =
            Com::CoCreateInstance(&Audio::MMDeviceEnumerator, None, Com::CLSCTX_ALL)
                .map_err(|e| format!("Failed to create device enumerator: {}", e))?;

        // Get default render (output) device - this is key for loopback!
        let device = enumerator
            .GetDefaultAudioEndpoint(Audio::eRender, Audio::eConsole)
            .map_err(|e| format!("Failed to get default render device: {}", e))?;

        // Get device name for logging
        if let Ok(id) = device.GetId() {
            info!("WASAPI loopback device: {:?}", id.to_string());
        }

        // Activate audio client
        let audio_client: Audio::IAudioClient = device
            .Activate(Com::CLSCTX_ALL, None)
            .map_err(|e| format!("Failed to activate audio client: {}", e))?;

        // Get the mix format (what the device is actually using)
        let mix_format = audio_client
            .GetMixFormat()
            .map_err(|e| format!("Failed to get mix format: {}", e))?;

        let format = &*mix_format;
        let sample_rate = format.nSamplesPerSec;
        let channels = format.nChannels;
        let bits_per_sample = format.wBitsPerSample;
        let block_align = format.nBlockAlign;

        info!(
            "WASAPI loopback format: {} Hz, {} ch, {} bits",
            sample_rate, channels, bits_per_sample
        );

        // Initialize audio client in loopback mode
        // Key flags: AUDCLNT_STREAMFLAGS_LOOPBACK for capturing output
        // Must use shared mode (not exclusive) for loopback
        let buffer_duration = 10_000_000i64; // 1 second in 100-nanosecond units

        audio_client
            .Initialize(
                Audio::AUDCLNT_SHAREMODE_SHARED,
                Audio::AUDCLNT_STREAMFLAGS_LOOPBACK,
                buffer_duration,
                0,
                mix_format,
                None,
            )
            .map_err(|e| format!("Failed to initialize audio client: {}", e))?;

        // Get capture client
        let capture_client: Audio::IAudioCaptureClient = audio_client
            .GetService()
            .map_err(|e| format!("Failed to get capture client: {}", e))?;

        // Start capturing
        audio_client
            .Start()
            .map_err(|e| format!("Failed to start audio client: {}", e))?;

        info!("WASAPI loopback capture started");

        // Capture loop - use polling since event mode doesn't work for loopback
        while !*should_stop.lock() {
            // Sleep a bit to avoid busy-waiting (10ms = 100Hz polling)
            thread::sleep(std::time::Duration::from_millis(10));

            // Get available frames
            let frames_available = match capture_client.GetNextPacketSize() {
                Ok(frames) => frames,
                Err(e) => {
                    warn!("Failed to get packet size: {}", e);
                    continue;
                }
            };

            if frames_available == 0 {
                continue;
            }

            // Get buffer
            let mut data_ptr: *mut u8 = std::ptr::null_mut();
            let mut num_frames: u32 = 0;
            let mut flags: u32 = 0;

            if let Err(e) = capture_client.GetBuffer(
                &mut data_ptr,
                &mut num_frames,
                &mut flags,
                None,
                None,
            ) {
                warn!("Failed to get buffer: {}", e);
                continue;
            }

            if num_frames > 0 && !data_ptr.is_null() {
                // Convert to f32 samples
                let samples = convert_to_f32(
                    data_ptr,
                    num_frames as usize,
                    channels as usize,
                    bits_per_sample,
                    block_align as usize,
                );

                // Convert to mono
                let mono_samples = to_mono(&samples, channels as usize);

                // Resample to 16kHz if needed
                let resampled = if sample_rate != TARGET_SAMPLE_RATE {
                    resample(&mono_samples, sample_rate, TARGET_SAMPLE_RATE)
                } else {
                    mono_samples
                };

                // Add to buffer
                buffer.lock().extend_from_slice(&resampled);
            }

            // Release buffer
            if let Err(e) = capture_client.ReleaseBuffer(num_frames) {
                warn!("Failed to release buffer: {}", e);
            }
        }

        // Stop and cleanup
        let _ = audio_client.Stop();

        Ok(())
    }
}

/// Convert raw audio bytes to f32 samples.
fn convert_to_f32(
    data: *mut u8,
    num_frames: usize,
    channels: usize,
    bits_per_sample: u16,
    block_align: usize,
) -> Vec<f32> {
    let total_samples = num_frames * channels;
    let mut samples = Vec::with_capacity(total_samples);

    unsafe {
        match bits_per_sample {
            16 => {
                let ptr = data as *const i16;
                for i in 0..total_samples {
                    let sample = *ptr.add(i);
                    samples.push(sample as f32 / 32768.0);
                }
            }
            32 => {
                // Could be f32 or i32 - WASAPI mix format is usually f32
                let ptr = data as *const f32;
                for i in 0..total_samples {
                    samples.push(*ptr.add(i));
                }
            }
            24 => {
                // 24-bit samples packed in 3 bytes
                for i in 0..total_samples {
                    let offset = (i / channels) * block_align + (i % channels) * 3;
                    let b0 = *data.add(offset) as i32;
                    let b1 = *data.add(offset + 1) as i32;
                    let b2 = *data.add(offset + 2) as i32;
                    let sample = (b2 << 16) | (b1 << 8) | b0;
                    // Sign extend from 24 to 32 bits
                    let sample = if sample & 0x800000 != 0 {
                        sample | 0xFF000000u32 as i32
                    } else {
                        sample
                    };
                    samples.push(sample as f32 / 8388608.0);
                }
            }
            _ => {
                warn!("Unsupported bits per sample: {}", bits_per_sample);
            }
        }
    }

    samples
}

/// Convert multi-channel audio to mono by averaging channels.
fn to_mono(samples: &[f32], channels: usize) -> Vec<f32> {
    if channels == 1 {
        return samples.to_vec();
    }

    samples
        .chunks(channels)
        .map(|chunk| chunk.iter().sum::<f32>() / channels as f32)
        .collect()
}

/// Simple linear interpolation resampling.
fn resample(samples: &[f32], from_rate: u32, to_rate: u32) -> Vec<f32> {
    if from_rate == to_rate || samples.is_empty() {
        return samples.to_vec();
    }

    let ratio = to_rate as f64 / from_rate as f64;
    let new_len = (samples.len() as f64 * ratio) as usize;
    let mut output = Vec::with_capacity(new_len);

    for i in 0..new_len {
        let src_idx = i as f64 / ratio;
        let src_idx_floor = src_idx.floor() as usize;
        let src_idx_ceil = (src_idx_floor + 1).min(samples.len() - 1);
        let frac = src_idx - src_idx_floor as f64;

        let sample =
            samples[src_idx_floor] as f64 * (1.0 - frac) + samples[src_idx_ceil] as f64 * frac;
        output.push(sample as f32);
    }

    output
}
