tenant.openagents/omega

1use std::{cell::RefCell, sync::mpsc::Receiver};
2
3use anyhow::{Context as _, Result};
4use windows::{
5    Win32::{
6        Foundation::{HWND, LPARAM, LRESULT, RECT, WPARAM},
7        Graphics::Gdi::{
8            BeginPaint, CLEARTYPE_QUALITY, CLIP_DEFAULT_PRECIS, CreateFontW, DEFAULT_CHARSET,
9            DeleteObject, EndPaint, FW_NORMAL, LOGFONTW, OUT_TT_ONLY_PRECIS, PAINTSTRUCT,
10            ReleaseDC, SelectObject, TextOutW,
11        },
12        System::LibraryLoader::GetModuleHandleW,
13        UI::{
14            Controls::{PBM_SETRANGE, PBM_SETSTEP, PBM_STEPIT, PROGRESS_CLASS},
15            WindowsAndMessaging::{
16                CREATESTRUCTW, CS_HREDRAW, CS_VREDRAW, CreateWindowExW, DefWindowProcW,
17                GWLP_USERDATA, GetDesktopWindow, GetWindowLongPtrW, GetWindowRect, HICON,
18                IMAGE_ICON, LR_DEFAULTSIZE, LR_SHARED, LoadImageW, PostQuitMessage, RegisterClassW,
19                SPI_GETICONTITLELOGFONT, SYSTEM_PARAMETERS_INFO_UPDATE_FLAGS, SendMessageW,
20                SetWindowLongPtrW, SystemParametersInfoW, WINDOW_EX_STYLE, WM_CLOSE, WM_CREATE,
21                WM_DESTROY, WM_NCCREATE, WM_PAINT, WNDCLASSW, WS_CAPTION, WS_CHILD, WS_EX_TOPMOST,
22                WS_POPUP, WS_VISIBLE,
23            },
24        },
25    },
26    core::HSTRING,
27};
28
29use crate::{
30    updater::JOBS,
31    windows_impl::{WM_JOB_UPDATED, WM_TERMINATE, show_error},
32};
33
34#[repr(C)]
35#[derive(Debug)]
36struct DialogInfo {
37    rx: Receiver<Result<()>>,
38    progress_bar: isize,
39}
40
41pub(crate) fn create_dialog_window(receiver: Receiver<Result<()>>) -> Result<HWND> {
42    unsafe {
43        let class_name = windows::core::w!("Zed-Auto-Updater-Dialog-Class");
44        let module = GetModuleHandleW(None).context("unable to get module handle")?;
45        let handle = LoadImageW(
46            Some(module.into()),
47            windows::core::PCWSTR(1 as _),
48            IMAGE_ICON,
49            0,
50            0,
51            LR_DEFAULTSIZE | LR_SHARED,
52        )
53        .context("unable to load icon file")?;
54        let wc = WNDCLASSW {
55            lpfnWndProc: Some(wnd_proc),
56            lpszClassName: class_name,
57            style: CS_HREDRAW | CS_VREDRAW,
58            hIcon: HICON(handle.0),
59            ..Default::default()
60        };
61        RegisterClassW(&wc);
62        let mut rect = RECT::default();
63        GetWindowRect(GetDesktopWindow(), &mut rect)
64            .context("unable to get desktop window rect")?;
65        let width = 400;
66        let height = 150;
67        let info = Box::new(RefCell::new(DialogInfo {
68            rx: receiver,
69            progress_bar: 0,
70        }));
71
72        let hwnd = CreateWindowExW(
73            WS_EX_TOPMOST,
74            class_name,
75            windows::core::w!("Omega"),
76            WS_VISIBLE | WS_POPUP | WS_CAPTION,
77            rect.right / 2 - width / 2,
78            rect.bottom / 2 - height / 2,
79            width,
80            height,
81            None,
82            None,
83            None,
84            Some(Box::into_raw(info) as _),
85        )
86        .context("unable to create dialog window")?;
87        Ok(hwnd)
88    }
89}
90
91macro_rules! return_if_failed {
92    ($e:expr) => {
93        match $e {
94            Ok(v) => v,
95            Err(e) => {
96                return LRESULT(e.code().0 as _);
97            }
98        }
99    };
100}
101
102macro_rules! make_lparam {
103    ($l:expr, $h:expr) => {
104        LPARAM(($l as u32 | ($h as u32) << 16) as isize)
105    };
106}
107
108unsafe extern "system" fn wnd_proc(
109    hwnd: HWND,
110    msg: u32,
111    wparam: WPARAM,
112    lparam: LPARAM,
113) -> LRESULT {
114    match msg {
115        WM_NCCREATE => unsafe {
116            let create_struct = lparam.0 as *const CREATESTRUCTW;
117            let info = (*create_struct).lpCreateParams as *mut RefCell<DialogInfo>;
118            let info = Box::from_raw(info);
119            SetWindowLongPtrW(hwnd, GWLP_USERDATA, Box::into_raw(info) as _);
120            DefWindowProcW(hwnd, msg, wparam, lparam)
121        },
122        WM_CREATE => unsafe {
123            // Create progress bar
124            let mut rect = RECT::default();
125            return_if_failed!(GetWindowRect(hwnd, &mut rect));
126            let progress_bar = return_if_failed!(CreateWindowExW(
127                WINDOW_EX_STYLE(0),
128                PROGRESS_CLASS,
129                None,
130                WS_CHILD | WS_VISIBLE,
131                20,
132                50,
133                340,
134                35,
135                Some(hwnd),
136                None,
137                None,
138                None,
139            ));
140            SendMessageW(
141                progress_bar,
142                PBM_SETRANGE,
143                None,
144                Some(make_lparam!(0, JOBS.len() * 10)),
145            );
146            SendMessageW(progress_bar, PBM_SETSTEP, Some(WPARAM(10)), None);
147            with_dialog_data(hwnd, |data| {
148                data.borrow_mut().progress_bar = progress_bar.0 as isize
149            });
150            LRESULT(0)
151        },
152        WM_PAINT => unsafe {
153            let mut ps = PAINTSTRUCT::default();
154            let hdc = BeginPaint(hwnd, &mut ps);
155
156            let font_name = get_system_ui_font_name();
157            let font = CreateFontW(
158                24,
159                0,
160                0,
161                0,
162                FW_NORMAL.0 as _,
163                0,
164                0,
165                0,
166                DEFAULT_CHARSET,
167                OUT_TT_ONLY_PRECIS,
168                CLIP_DEFAULT_PRECIS,
169                CLEARTYPE_QUALITY,
170                0,
171                &HSTRING::from(font_name),
172            );
173            let temp = SelectObject(hdc, font.into());
174            let string = HSTRING::from("Updating Omega...");
175            return_if_failed!(TextOutW(hdc, 20, 15, &string).ok());
176            return_if_failed!(DeleteObject(temp).ok());
177
178            return_if_failed!(EndPaint(hwnd, &ps).ok());
179            ReleaseDC(Some(hwnd), hdc);
180
181            LRESULT(0)
182        },
183        WM_JOB_UPDATED => with_dialog_data(hwnd, |data| {
184            let progress_bar = data.borrow().progress_bar;
185            unsafe { SendMessageW(HWND(progress_bar as _), PBM_STEPIT, None, None) }
186        }),
187        WM_TERMINATE => {
188            with_dialog_data(hwnd, |data| {
189                if let Ok(result) = data.borrow_mut().rx.recv()
190                    && let Err(e) = result
191                {
192                    log::error!("Failed to update Omega: {:?}", e);
193                    show_error(format!("Error: {:?}", e));
194                }
195            });
196            unsafe { PostQuitMessage(0) };
197            LRESULT(0)
198        }
199        WM_CLOSE => LRESULT(0), // Prevent user occasionally closing the window
200        WM_DESTROY => {
201            unsafe { PostQuitMessage(0) };
202            LRESULT(0)
203        }
204        _ => unsafe { DefWindowProcW(hwnd, msg, wparam, lparam) },
205    }
206}
207
208fn with_dialog_data<F, T>(hwnd: HWND, f: F) -> T
209where
210    F: FnOnce(&RefCell<DialogInfo>) -> T,
211{
212    let raw = unsafe { GetWindowLongPtrW(hwnd, GWLP_USERDATA) as *mut RefCell<DialogInfo> };
213    let data = unsafe { Box::from_raw(raw) };
214    let result = f(data.as_ref());
215    unsafe { SetWindowLongPtrW(hwnd, GWLP_USERDATA, Box::into_raw(data) as _) };
216    result
217}
218
219fn get_system_ui_font_name() -> String {
220    unsafe {
221        let mut info: LOGFONTW = std::mem::zeroed();
222        if SystemParametersInfoW(
223            SPI_GETICONTITLELOGFONT,
224            std::mem::size_of::<LOGFONTW>() as u32,
225            Some(&mut info as *mut _ as _),
226            SYSTEM_PARAMETERS_INFO_UPDATE_FLAGS(0),
227        )
228        .is_ok()
229        {
230            let font_name = String::from_utf16_lossy(&info.lfFaceName);
231            font_name.trim_matches(char::from(0)).to_owned()
232        } else {
233            "MS Shell Dlg".to_owned()
234        }
235    }
236}
237