KiCad - the KiCad Bridge

//! KiCad's OWN AppUserModelID onto KiCad's OWN standalone Start Menu shortcuts.
//!
//! The white taskbar icon, solved end to end on John's taskbar on 2026-08-23 and recorded
//! in skills/kicad-bridge-dev/SKILL.md ("THE WHITE-ICON MECHANISM"): a standalone editor
//! process, which is how this bridge launches pcbnew, presents a runtime-built AUMID such
//! as `Kicad.Kicad.pcbnew.10.0`. KiCad's installer writes that AUMID onto the MAIN
//! shortcut only and leaves it blank on the standalone editor shortcuts, so the shell has
//! no shortcut to resolve the editor group's icon from. Some machines fall back to the
//! window's own icon; others paint the generic white page. Writing KiCad's own AUMID onto
//! KiCad's own user-level shortcut makes the resolution deterministic, the exact mechanism
//! that keeps the main KiCad button correct. No Adom identity is involved anywhere: this
//! bridge stamps NO AUMIDs of its own, ever.
//!
//! The Python bridge did this through inline C# in PowerShell (`handlers/win_taskbar.py`,
//! `_repair_kicad_shortcut_aumids`); this is the same operation in process. Idempotent: a
//! shortcut that already carries any AUMID is left alone. The shell reads the mapping when
//! it next rebuilds its Start Menu index (logon or an explorer restart); `SHChangeNotify`
//! is sent for each written file to invite an earlier refresh, and explorer is never
//! restarted from here.

use std::path::{Path, PathBuf};

use windows::core::{Interface, GUID, HSTRING, PWSTR};
use windows::Win32::Foundation::{PROPERTYKEY, RPC_E_CHANGED_MODE};
use windows::Win32::System::Com::StructuredStorage::{PropVariantClear, PROPVARIANT};
use windows::Win32::System::Com::{
    CoCreateInstance, CoInitializeEx, CoTaskMemAlloc, CoTaskMemFree, CoUninitialize, IPersistFile,
    CLSCTX_INPROC_SERVER, COINIT_APARTMENTTHREADED, STGM_READWRITE,
};
use windows::Win32::System::Variant::VT_LPWSTR;
use windows::Win32::UI::Shell::PropertiesSystem::IPropertyStore;
use windows::Win32::UI::Shell::{IShellLinkW, SHChangeNotify, ShellLink, SHCNE_UPDATEITEM, SHCNF_FLUSHNOWAIT, SHCNF_PATHW};

/// `System.AppUserModel.ID` (propkey.h): {9F4C2855-9F79-4B39-A8D0-E1D42DE1D5F3}, 5.
const PKEY_APPUSERMODEL_ID: PROPERTYKEY = PROPERTYKEY { fmtid: GUID::from_u128(0x9F4C2855_9F79_4B39_A8D0_E1D42DE1D5F3), pid: 5 };

/// One shortcut's fate.
#[derive(Debug, Clone)]
pub struct ShortcutRepair {
    pub path: String,
    pub aumid: String,
    /// `written`, `already` (carried an AUMID, untouched), `missing` (no such shortcut), `error`.
    pub outcome: &'static str,
    pub detail: Option<String>,
}

/// The standalone shortcuts KiCad's installer writes without an AUMID, with the AUMID each
/// running process actually presents (read from the taskbar buttons' AutomationIds).
/// `version` is major.minor ("10.0"): KiCad names both the folder and the shortcuts by it.
pub fn shortcuts_for(version: &str) -> Vec<(PathBuf, String)> {
    // A per-user install puts its Start Menu folder under %APPDATA%; an all-users install (what
    // kicad_upgrade does when the bridge runs elevated) puts it under %ProgramData%. Looking only
    // at the first reported all three shortcuts missing on a fresh all-users box (winvm, 10.0.6).
    let folder = |base: std::ffi::OsString| Path::new(&base).join("Microsoft").join("Windows").join("Start Menu").join("Programs").join(format!("KiCad {version}"));
    let roots: Vec<PathBuf> = ["APPDATA", "ProgramData"].iter().filter_map(|v| std::env::var_os(v)).map(folder).collect();
    if roots.is_empty() {
        return Vec::new();
    }
    [
        (format!("PCB Editor {version} (standalone).lnk"), format!("Kicad.Kicad.pcbnew.{version}")),
        (format!("Schematic Editor {version} (standalone).lnk"), format!("Kicad.Kicad.eeschema.{version}")),
        (format!("Gerber Viewer {version}.lnk"), format!("Kicad.Kicad.gerbview.{version}")),
    ]
    .into_iter()
    .map(|(name, aumid)| {
        let found = roots.iter().map(|r| r.join(&name)).find(|p| p.is_file());
        (found.unwrap_or_else(|| roots[0].join(&name)), aumid)
    })
    .collect()
}

/// COM apartment for this thread, released on drop; a thread already in an apartment of
/// the other kind keeps it (RPC_E_CHANGED_MODE is not a failure for us).
struct Com(bool);

impl Com {
    fn init() -> Result<Com, String> {
        // SAFETY: CoInitializeEx on the current thread; balanced in Drop only when this
        // call took the reference.
        let hr = unsafe { CoInitializeEx(None, COINIT_APARTMENTTHREADED) };
        if hr.is_ok() {
            Ok(Com(true))
        } else if hr == RPC_E_CHANGED_MODE {
            Ok(Com(false))
        } else {
            Err(format!("CoInitializeEx: {hr}"))
        }
    }
}

impl Drop for Com {
    fn drop(&mut self) {
        if self.0 {
            // SAFETY: matches the successful CoInitializeEx; every COM object made in
            // between is dropped before this runs (they live inside repair_one).
            unsafe { CoUninitialize() };
        }
    }
}

/// Fill the AUMID on every standalone shortcut that lacks one. Never fails as a whole:
/// each shortcut reports its own outcome, and a missing shortcut is `missing`, not an error.
pub fn repair_kicad_shortcut_aumids(version: &str) -> Result<Vec<ShortcutRepair>, String> {
    let _com = Com::init()?;
    let mut out = Vec::new();
    for (path, aumid) in shortcuts_for(version) {
        let shown = path.to_string_lossy().replace('\\', "/");
        if !path.is_file() {
            out.push(ShortcutRepair { path: shown, aumid, outcome: "missing", detail: None });
            continue;
        }
        match repair_one(&path, &aumid) {
            Ok(true) => out.push(ShortcutRepair { path: shown, aumid, outcome: "written", detail: None }),
            Ok(false) => out.push(ShortcutRepair { path: shown, aumid, outcome: "already", detail: None }),
            Err(e) => out.push(ShortcutRepair { path: shown, aumid, outcome: "error", detail: Some(e) }),
        }
    }
    Ok(out)
}

/// Ok(true) written, Ok(false) already carried an AUMID.
fn repair_one(path: &Path, aumid: &str) -> Result<bool, String> {
    let wide = HSTRING::from(path.as_os_str());
    // SAFETY: standard COM object creation and QueryInterface on a live object; every
    // pointer handed to the shell outlives the call that reads it.
    unsafe {
        let link: IShellLinkW = CoCreateInstance(&ShellLink, None, CLSCTX_INPROC_SERVER).map_err(|e| format!("CoCreateInstance(ShellLink): {e}"))?;
        let file: IPersistFile = link.cast().map_err(|e| format!("IPersistFile: {e}"))?;
        file.Load(&wide, STGM_READWRITE).map_err(|e| format!("Load: {e}"))?;
        let store: IPropertyStore = link.cast().map_err(|e| format!("IPropertyStore: {e}"))?;
        let mut current = store.GetValue(&PKEY_APPUSERMODEL_ID).map_err(|e| format!("GetValue: {e}"))?;
        let existing = read_lpwstr(&current);
        let _ = PropVariantClear(&mut current);
        if existing.map(|s| !s.is_empty()).unwrap_or(false) {
            return Ok(false);
        }
        let mut value = lpwstr_variant(aumid);
        let set = store.SetValue(&PKEY_APPUSERMODEL_ID, &value).map_err(|e| format!("SetValue: {e}"));
        let commit = set.and_then(|_| store.Commit().map_err(|e| format!("Commit: {e}")));
        free_lpwstr_variant(&mut value);
        commit?;
        file.Save(&wide, true).map_err(|e| format!("Save: {e}"))?;
        // Invite explorer to refresh its view of this item now rather than at next logon.
        SHChangeNotify(SHCNE_UPDATEITEM, SHCNF_PATHW | SHCNF_FLUSHNOWAIT, Some(wide.as_ptr() as *const _), None);
    }
    Ok(true)
}

/// The string in a VT_LPWSTR PROPVARIANT, None for any other type (VT_EMPTY on a shortcut
/// that has no AUMID).
unsafe fn read_lpwstr(pv: &PROPVARIANT) -> Option<String> {
    let inner = &pv.Anonymous.Anonymous;
    if inner.vt != VT_LPWSTR {
        return None;
    }
    let p: PWSTR = inner.Anonymous.pwszVal;
    if p.is_null() {
        return Some(String::new());
    }
    Some(p.to_string().unwrap_or_default())
}

/// A VT_LPWSTR PROPVARIANT owning a CoTaskMem copy of `s` (what IPropertyStore expects).
unsafe fn lpwstr_variant(s: &str) -> PROPVARIANT {
    let mut units: Vec<u16> = s.encode_utf16().collect();
    units.push(0);
    let bytes = units.len() * 2;
    let mem = CoTaskMemAlloc(bytes) as *mut u16;
    std::ptr::copy_nonoverlapping(units.as_ptr(), mem, units.len());
    let mut pv = PROPVARIANT::default();
    let inner = &mut *(&mut pv.Anonymous.Anonymous as *mut _ as *mut windows::Win32::System::Com::StructuredStorage::PROPVARIANT_0_0);
    inner.vt = VT_LPWSTR;
    inner.Anonymous.pwszVal = PWSTR(mem);
    pv
}

unsafe fn free_lpwstr_variant(pv: &mut PROPVARIANT) {
    let inner = &pv.Anonymous.Anonymous;
    let p = inner.Anonymous.pwszVal;
    if !p.is_null() {
        CoTaskMemFree(Some(p.0 as *const _));
    }
}