// File: src/audio.rs
// Author: Hadi Cahyadi <cumulus13@gmail.com>
// Windows Audio Session API (WASAPI) backend via windows-rs

use anyhow::{anyhow, Result};
use windows::{
    core::{Interface, GUID},
    Win32::{
        Foundation::S_FALSE,
        Media::Audio::{
            eConsole, eRender,
            IAudioSessionControl, IAudioSessionControl2, IAudioSessionEnumerator,
            IAudioSessionManager2, IMMDevice, IMMDeviceEnumerator, ISimpleAudioVolume,
            MMDeviceEnumerator,
        },
        Media::Audio::Endpoints::IAudioEndpointVolume,
        System::Com::{CoCreateInstance, CoInitializeEx, CLSCTX_ALL, COINIT_APARTMENTTHREADED},
        System::Threading::{
            OpenProcess, QueryFullProcessImageNameW,
            PROCESS_NAME_WIN32, PROCESS_QUERY_LIMITED_INFORMATION,
        },
    },
};
use std::collections::HashMap;
use std::ffi::OsString;
use std::os::windows::ffi::OsStringExt;

// ── COM init ──────────────────────────────────────────────────────────────────

pub fn init_com() -> Result<()> {
    unsafe {
        let hr = CoInitializeEx(None, COINIT_APARTMENTTHREADED);
        // S_FALSE means COM was already initialized on this thread — that is fine.
        if hr.is_err() && hr != S_FALSE {
            return Err(anyhow!("CoInitializeEx failed: {:?}", hr));
        }
    }
    Ok(())
}

// ── Public types ──────────────────────────────────────────────────────────────

#[derive(Debug, Clone, PartialEq)]
pub enum ChannelKind {
    Master,
    /// Carries the process-id so mutations can target the exact session(s).
    Application { pid: u32 },
}

#[derive(Debug, Clone)]
pub struct AudioChannel {
    pub id: usize,
    pub name: String,
    pub volume: u8, // 0-100
    pub muted: bool,
    pub kind: ChannelKind,
}

// ── WASAPI mixer ──────────────────────────────────────────────────────────────

pub struct WasapiMixer {
    endpoint_volume: IAudioEndpointVolume,
    session_manager: IAudioSessionManager2,
}

impl WasapiMixer {
    pub fn new() -> Result<Self> {
        unsafe {
            let enumerator: IMMDeviceEnumerator =
                CoCreateInstance(&MMDeviceEnumerator, None, CLSCTX_ALL)?;
            let device: IMMDevice =
                enumerator.GetDefaultAudioEndpoint(eRender, eConsole)?;
            let endpoint_volume: IAudioEndpointVolume = device.Activate(CLSCTX_ALL, None)?;
            let session_manager: IAudioSessionManager2 = device.Activate(CLSCTX_ALL, None)?;
            Ok(Self { endpoint_volume, session_manager })
        }
    }

    // ── Master volume ─────────────────────────────────────────────────────

    fn master_volume(&self) -> Result<(u8, bool)> {
        unsafe {
            let scalar: f32 = self.endpoint_volume.GetMasterVolumeLevelScalar()?;
            let muted = self.endpoint_volume.GetMute()?.as_bool();
            Ok(((scalar * 100.0).round() as u8, muted))
        }
    }

    fn set_master_scalar(&self, scalar: f32) -> Result<()> {
        unsafe {
            self.endpoint_volume
                .SetMasterVolumeLevelScalar(scalar.clamp(0.0, 1.0), &GUID::zeroed())?
        }
        Ok(())
    }

    fn set_master_mute(&self, muted: bool) -> Result<()> {
        unsafe { self.endpoint_volume.SetMute(muted, &GUID::zeroed())? }
        Ok(())
    }

    // ── Application sessions ──────────────────────────────────────────────
    //
    // How we identify app sessions vs. system/ghost sessions:
    //
    //   The "System Sounds" pseudo-session has PID 0 — OpenProcess() on pid 0
    //   fails and process_name() returns "". Every real application session has
    //   a non-zero PID that resolves to an executable name.  This is exactly
    //   the same strategy used by pycaw (Python): filter by `session.Process`
    //   being non-None, which maps directly to pid != 0 here.
    //
    //   We intentionally do NOT call IsSystemSoundsSession() because the
    //   windows-rs crate collapses both S_OK and S_FALSE into Ok(()), making
    //   the return value unreadable without unsafe vtable tricks that are
    //   fragile across crate versions. The pid == 0 check is simpler, safer,
    //   and correct.
    //
    // Additional filtering:
    //   - Skip sessions where process_name() returns "" (access denied /
    //     process already exited but session not yet expired).
    //   - Group multiple WASAPI streams from the same PID into one channel
    //     (e.g. browsers open several streams). Volume = average, mute = OR.
    //   - Sort by name so the list is stable across refreshes.

    fn app_sessions(&self) -> Result<Vec<(u32, String, u8, bool)>> {
        // pid → (display_name, volume_sum, stream_count, any_muted)
        let mut by_pid: HashMap<u32, (String, u32, u32, bool)> = HashMap::new();

        unsafe {
            let enumerator: IAudioSessionEnumerator =
                self.session_manager.GetSessionEnumerator()?;
            let count = enumerator.GetCount()?;

            for i in 0..count {
                // ── Get session control interfaces ────────────────────
                let ctl: IAudioSessionControl = match enumerator.GetSession(i) {
                    Ok(c) => c,
                    Err(_) => continue,
                };
                let ctl2: IAudioSessionControl2 = match ctl.cast() {
                    Ok(c) => c,
                    Err(_) => continue,
                };

                // ── Filter: must have a real PID ──────────────────────
                // pid == 0  →  System Sounds pseudo-session or ghost session
                let pid = match ctl2.GetProcessId() {
                    Ok(p) if p != 0 => p,
                    _ => continue,
                };

                // ── Filter: PID must resolve to an executable name ────
                // Fails for terminated processes whose sessions haven't
                // been collected yet, and for protected system processes.
                let name = process_name(pid);
                if name.is_empty() {
                    continue;
                }

                // ── Read per-stream volume & mute ─────────────────────
                let sav: ISimpleAudioVolume = match ctl.cast() {
                    Ok(s) => s,
                    Err(_) => continue,
                };
                let vol_f: f32 = match sav.GetMasterVolume() {
                    Ok(v) => v,
                    Err(_) => continue,
                };
                let muted = sav.GetMute().map(|b| b.as_bool()).unwrap_or(false);
                let vol_pct = (vol_f * 100.0).round() as u32;

                // ── Merge into per-PID bucket ─────────────────────────
                let entry = by_pid.entry(pid).or_insert((name, 0, 0, false));
                entry.1 += vol_pct;  // accumulate volume
                entry.2 += 1;        // count streams
                entry.3 |= muted;    // any stream muted → channel muted
            }
        }

        // Flatten: average volume per PID, then sort by name
        let mut result: Vec<(u32, String, u8, bool)> = by_pid
            .into_iter()
            .map(|(pid, (name, sum, count, muted))| {
                let avg = if count > 0 { (sum / count).min(100) as u8 } else { 0 };
                (pid, name, avg, muted)
            })
            .collect();

        result.sort_by(|a, b| a.1.to_lowercase().cmp(&b.1.to_lowercase()));
        Ok(result)
    }

    // ── Public API ────────────────────────────────────────────────────────

    /// Full snapshot: Master first, then one entry per active application.
    pub fn snapshot(&self) -> Result<Vec<AudioChannel>> {
        let mut channels = Vec::new();
        let (master_vol, master_mute) = self.master_volume()?;

        channels.push(AudioChannel {
            id: 0,
            name: "Master".to_string(),
            volume: master_vol,
            muted: master_mute,
            kind: ChannelKind::Master,
        });

        for (idx, (pid, name, vol, muted)) in self.app_sessions()?.into_iter().enumerate() {
            channels.push(AudioChannel {
                id: idx + 1,
                name,
                volume: vol,
                muted,
                kind: ChannelKind::Application { pid },
            });
        }
        Ok(channels)
    }

    /// Adjust volume by `delta` percent. Re-reads live value so rapid
    /// key-repeats accumulate correctly rather than replaying the same base.
    pub fn adjust_volume(&self, channel: &AudioChannel, delta: i8) -> Result<()> {
        match channel.kind {
            ChannelKind::Master => {
                let live = unsafe { self.endpoint_volume.GetMasterVolumeLevelScalar()? };
                self.set_master_scalar((live + delta as f32 / 100.0).clamp(0.0, 1.0))
            }
            ChannelKind::Application { pid } => self.adjust_session_volume(pid, delta),
        }
    }

    /// Toggle mute on the given channel.
    pub fn toggle_mute(&self, channel: &AudioChannel) -> Result<()> {
        match channel.kind {
            ChannelKind::Master => self.set_master_mute(!channel.muted),
            ChannelKind::Application { pid } => self.set_session_mute(pid, !channel.muted),
        }
    }

    // ── Session mutation helpers ──────────────────────────────────────────

    fn adjust_session_volume(&self, target_pid: u32, delta: i8) -> Result<()> {
        self.for_each_session_by_pid(target_pid, |sav| unsafe {
            if let Ok(live) = sav.GetMasterVolume() {
                let _ = sav.SetMasterVolume(
                    (live + delta as f32 / 100.0).clamp(0.0, 1.0),
                    &GUID::zeroed(),
                );
            }
        })
    }

    fn set_session_mute(&self, target_pid: u32, muted: bool) -> Result<()> {
        self.for_each_session_by_pid(target_pid, |sav| unsafe {
            let _ = sav.SetMute(muted, &GUID::zeroed());
        })
    }

    /// Iterate every WASAPI stream for `target_pid` and call `f` on each.
    fn for_each_session_by_pid<F>(&self, target_pid: u32, f: F) -> Result<()>
    where
        F: Fn(&ISimpleAudioVolume),
    {
        unsafe {
            let enumerator: IAudioSessionEnumerator =
                self.session_manager.GetSessionEnumerator()?;
            let count = enumerator.GetCount()?;

            for i in 0..count {
                let ctl: IAudioSessionControl = match enumerator.GetSession(i) {
                    Ok(c) => c,
                    Err(_) => continue,
                };
                let ctl2: IAudioSessionControl2 = match ctl.cast() {
                    Ok(c) => c,
                    Err(_) => continue,
                };
                if ctl2.GetProcessId().unwrap_or(0) == target_pid {
                    if let Ok(sav) = ctl.cast::<ISimpleAudioVolume>() {
                        f(&sav);
                    }
                }
            }
        }
        Ok(())
    }
}

// ── Process name helper ───────────────────────────────────────────────────────

/// Return the executable stem (e.g. "chrome") for `pid`, or "" on any failure.
fn process_name(pid: u32) -> String {
    if pid == 0 {
        return String::new();
    }
    unsafe {
        let handle = match OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, false, pid) {
            Ok(h) if !h.is_invalid() => h,
            _ => return String::new(),
        };

        let mut buf = vec![0u16; 1024];
        let mut size = buf.len() as u32;

        if QueryFullProcessImageNameW(
            handle,
            PROCESS_NAME_WIN32,
            windows::core::PWSTR(buf.as_mut_ptr()),
            &mut size,
        )
        .is_ok()
        {
            let path = OsString::from_wide(&buf[..size as usize]);
            std::path::Path::new(path.to_string_lossy().as_ref())
                .file_stem()
                .and_then(|s| s.to_str())
                .unwrap_or("")
                .to_string()
        } else {
            String::new()
        }
    }
}
