headless_engine/dom/
screenshot.rs

1use crate::dom::interactive::InteractiveElement;
2use crate::dom::SearchResults;
3use anyhow::Result;
4use base64::Engine;
5use scraper::{Html, Selector};
6use serde::{Deserialize, Serialize};
7use std::path::PathBuf;
8
9use std::sync::atomic::{AtomicUsize, Ordering};
10
11#[derive(Debug, Clone, Serialize, Deserialize)]
12pub struct ScreenshotResult {
13    pub width: u32,
14    pub height: u32,
15    pub svg: String,
16    pub layout_wireframe: String,
17    pub element_count: usize,
18    #[serde(skip_serializing_if = "Vec::is_empty")]
19    pub png_bytes: Vec<u8>,
20    pub png_base64: String,
21}
22
23pub struct RealBrowserScreenshot;
24
25struct TempFileCleanup(PathBuf);
26impl Drop for TempFileCleanup {
27    fn drop(&mut self) {
28        let _ = std::fs::remove_file(&self.0);
29    }
30}
31
32impl RealBrowserScreenshot {
33    pub async fn dump_rendered_dom(url: &str) -> Option<String> {
34        let browser_bin = Self::find_browser_binary()?;
35        let temp_dir = std::env::temp_dir();
36        let pid = std::process::id();
37        let time = std::time::SystemTime::now()
38            .duration_since(std::time::UNIX_EPOCH)
39            .map(|d| d.as_millis())
40            .unwrap_or(0);
41        static DUMP_COUNTER: AtomicUsize = AtomicUsize::new(0);
42        let count = DUMP_COUNTER.fetch_add(1, Ordering::Relaxed);
43        let profile_dir = temp_dir.join(format!("headless_engine_dump_{}_{}_{}", pid, time, count));
44        let profile_arg = format!("--user-data-dir={}", profile_dir.to_string_lossy());
45
46        let mut cmd = tokio::process::Command::new(&browser_bin);
47        cmd.arg("--headless=new")
48            .arg("--disable-gpu")
49            .arg("--no-sandbox")
50            .arg(&profile_arg)
51            .arg("--hide-scrollbars")
52            .arg("--run-all-compositor-stages-before-draw")
53            .arg("--disable-blink-features=AutomationControlled")
54            .arg("--user-agent=Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/133.0.0.0 Safari/537.36")
55            .arg("--lang=en-US,en")
56            .arg("--dump-dom")
57            .arg(url);
58
59        #[cfg(target_os = "windows")]
60        cmd.creation_flags(0x08000000); // CREATE_NO_WINDOW
61
62        let output = tokio::time::timeout(std::time::Duration::from_secs(12), cmd.output())
63            .await
64            .ok()?
65            .ok()?;
66        let _ = std::fs::remove_dir_all(&profile_dir);
67
68        if output.status.success() && !output.stdout.is_empty() {
69            let html = String::from_utf8_lossy(&output.stdout).to_string();
70            if !html.is_empty() {
71                return Some(html);
72            }
73        }
74        None
75    }
76
77    pub fn find_browser_binary() -> Option<PathBuf> {
78        #[cfg(target_os = "windows")]
79        {
80            let candidates = [
81                r"C:\Program Files\Google\Chrome\Application\chrome.exe",
82                r"C:\Program Files (x86)\Google\Chrome\Application\chrome.exe",
83                r"C:\Program Files (x86)\Microsoft\Edge\Application\msedge.exe",
84                r"C:\Program Files\Microsoft\Edge\Application\msedge.exe",
85            ];
86            for path_str in &candidates {
87                let p = std::path::Path::new(path_str);
88                if p.exists() {
89                    return Some(p.to_path_buf());
90                }
91            }
92            if let Ok(local_app_data) = std::env::var("LOCALAPPDATA") {
93                let chrome_local = std::path::Path::new(&local_app_data)
94                    .join(r"Google\Chrome\Application\chrome.exe");
95                if chrome_local.exists() {
96                    return Some(chrome_local);
97                }
98                let edge_local = std::path::Path::new(&local_app_data)
99                    .join(r"Microsoft\Edge\Application\msedge.exe");
100                if edge_local.exists() {
101                    return Some(edge_local);
102                }
103            }
104            // Dynamic fallback
105            if let Ok(output) = std::process::Command::new("where").arg("chrome").output() {
106                if output.status.success() {
107                    let path_str = String::from_utf8_lossy(&output.stdout)
108                        .lines()
109                        .next()
110                        .unwrap_or("")
111                        .trim()
112                        .to_string();
113                    if !path_str.is_empty() {
114                        return Some(PathBuf::from(path_str));
115                    }
116                }
117            }
118            if let Ok(output) = std::process::Command::new("where").arg("msedge").output() {
119                if output.status.success() {
120                    let path_str = String::from_utf8_lossy(&output.stdout)
121                        .lines()
122                        .next()
123                        .unwrap_or("")
124                        .trim()
125                        .to_string();
126                    if !path_str.is_empty() {
127                        return Some(PathBuf::from(path_str));
128                    }
129                }
130            }
131        }
132
133        #[cfg(target_os = "macos")]
134        {
135            let candidates = [
136                "/Applications/Google Chrome.app/Contents/MacOS/Google Chrome",
137                "/Applications/Microsoft Edge.app/Contents/MacOS/Microsoft Edge",
138                "/Applications/Chromium.app/Contents/MacOS/Chromium",
139            ];
140            for path_str in &candidates {
141                let p = std::path::Path::new(path_str);
142                if p.exists() {
143                    return Some(p.to_path_buf());
144                }
145            }
146            // Dynamic fallback
147            if let Ok(output) = std::process::Command::new("which")
148                .arg("google-chrome")
149                .output()
150            {
151                if output.status.success() {
152                    let path_str = String::from_utf8_lossy(&output.stdout).trim().to_string();
153                    if !path_str.is_empty() {
154                        return Some(PathBuf::from(path_str));
155                    }
156                }
157            }
158        }
159
160        #[cfg(target_os = "linux")]
161        {
162            let candidates = [
163                "google-chrome",
164                "google-chrome-stable",
165                "chromium",
166                "chromium-browser",
167            ];
168            for bin in &candidates {
169                if let Ok(output) = std::process::Command::new("which").arg(bin).output() {
170                    if output.status.success() {
171                        let path_str = String::from_utf8_lossy(&output.stdout).trim().to_string();
172                        if !path_str.is_empty() {
173                            return Some(PathBuf::from(path_str));
174                        }
175                    }
176                }
177            }
178        }
179
180        None
181    }
182
183    fn generate_temp_png_path() -> PathBuf {
184        let temp_dir = std::env::temp_dir();
185        let pid = std::process::id();
186        let time = std::time::SystemTime::now()
187            .duration_since(std::time::UNIX_EPOCH)
188            .map(|d| d.as_millis())
189            .unwrap_or(0);
190        static COUNTER: AtomicUsize = AtomicUsize::new(0);
191        let count = COUNTER.fetch_add(1, Ordering::Relaxed);
192        temp_dir.join(format!("headless_shot_{}_{}_{}.png", pid, time, count))
193    }
194
195    fn is_valid_url(url: &str) -> bool {
196        url.starts_with("http://") || url.starts_with("https://") || url.starts_with("file://")
197    }
198
199    pub async fn capture_real_screenshot_async(
200        url: &str,
201        html_str: &str,
202        width: u32,
203        height: u32,
204    ) -> Option<ScreenshotResult> {
205        if !Self::is_valid_url(url) {
206            return None;
207        }
208
209        let browser_bin = Self::find_browser_binary()?;
210
211        let temp_png = Self::generate_temp_png_path();
212        let mut temp_html = temp_png.clone();
213        temp_html.set_extension("html");
214
215        let mut injected_html = html_str.to_string();
216        let base_tag = format!("<base href=\"{}\">", url);
217        if let Some(idx) = injected_html.find("<head>") {
218            injected_html.insert_str(idx + 6, &base_tag);
219        } else if let Some(idx) = injected_html.find("<head ") {
220            if let Some(close_idx) = injected_html[idx..].find('>') {
221                injected_html.insert_str(idx + close_idx + 1, &base_tag);
222            }
223        } else {
224            injected_html.insert_str(0, &base_tag);
225        }
226
227        let _ = std::fs::write(&temp_html, injected_html);
228
229        let _cleanup_png = TempFileCleanup(temp_png.clone());
230        let _cleanup_html = TempFileCleanup(temp_html.clone());
231
232        let temp_png_str = temp_png.to_string_lossy().to_string();
233        let temp_html_str = format!(
234            "file:///{}",
235            temp_html.to_string_lossy().to_string().replace('\\', "/")
236        );
237
238        let screenshot_arg = format!("--screenshot={}", temp_png_str);
239        let window_size_arg = format!("--window-size={},{}", width, height);
240
241        let target_arg = if url.starts_with("http://") || url.starts_with("https://") {
242            url.to_string()
243        } else {
244            temp_html_str
245        };
246
247        let profile_dir = temp_png.with_extension("profile");
248        let profile_arg = format!("--user-data-dir={}", profile_dir.to_string_lossy());
249
250        let mut cmd = tokio::process::Command::new(&browser_bin);
251        cmd.arg("--headless=new")
252            .arg("--disable-gpu")
253            .arg("--no-sandbox")
254            .arg(&profile_arg)
255            .arg("--hide-scrollbars")
256            .arg("--run-all-compositor-stages-before-draw")
257            .arg("--disable-blink-features=AutomationControlled")
258            .arg("--user-agent=Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/133.0.0.0 Safari/537.36")
259            .arg("--lang=en-US,en")
260            .arg(&window_size_arg)
261            .arg(&screenshot_arg)
262            .arg(&target_arg);
263
264        #[cfg(target_os = "windows")]
265        cmd.creation_flags(0x08000000); // CREATE_NO_WINDOW
266
267        if let Ok(child) = cmd.spawn() {
268            let _ =
269                tokio::time::timeout(std::time::Duration::from_secs(12), child.wait_with_output())
270                    .await;
271        }
272
273        let _ = std::fs::remove_dir_all(&profile_dir);
274
275        if temp_png.exists() {
276            if let Ok(png_bytes) = std::fs::read(&temp_png) {
277                let png_base64 = format!(
278                    "data:image/png;base64,{}",
279                    base64::engine::general_purpose::STANDARD.encode(&png_bytes)
280                );
281                return Some(ScreenshotResult {
282                    width,
283                    height,
284                    svg: String::new(),
285                    layout_wireframe: String::new(),
286                    element_count: 1,
287                    png_bytes,
288                    png_base64,
289                });
290            }
291        }
292        None
293    }
294
295    pub fn capture_real_screenshot_sync(
296        url: &str,
297        html_str: &str,
298        width: u32,
299        height: u32,
300    ) -> Option<ScreenshotResult> {
301        if !Self::is_valid_url(url) {
302            return None;
303        }
304
305        let browser_bin = Self::find_browser_binary()?;
306
307        let temp_png = Self::generate_temp_png_path();
308        let mut temp_html = temp_png.clone();
309        temp_html.set_extension("html");
310
311        let mut injected_html = html_str.to_string();
312        let base_tag = format!("<base href=\"{}\">", url);
313        if let Some(idx) = injected_html.find("<head>") {
314            injected_html.insert_str(idx + 6, &base_tag);
315        } else if let Some(idx) = injected_html.find("<head ") {
316            if let Some(close_idx) = injected_html[idx..].find('>') {
317                injected_html.insert_str(idx + close_idx + 1, &base_tag);
318            }
319        } else {
320            injected_html.insert_str(0, &base_tag);
321        }
322
323        let _ = std::fs::write(&temp_html, injected_html);
324
325        let _cleanup_png = TempFileCleanup(temp_png.clone());
326        let _cleanup_html = TempFileCleanup(temp_html.clone());
327
328        let temp_png_str = temp_png.to_string_lossy().to_string();
329        let temp_html_str = format!(
330            "file:///{}",
331            temp_html.to_string_lossy().to_string().replace('\\', "/")
332        );
333
334        let screenshot_arg = format!("--screenshot={}", temp_png_str);
335        let window_size_arg = format!("--window-size={},{}", width, height);
336
337        let target_arg = if url.starts_with("http://") || url.starts_with("https://") {
338            url.to_string()
339        } else {
340            temp_html_str
341        };
342
343        let profile_dir = temp_png.with_extension("profile");
344        let profile_arg = format!("--user-data-dir={}", profile_dir.to_string_lossy());
345
346        let mut cmd = std::process::Command::new(&browser_bin);
347        cmd.arg("--headless=new")
348            .arg("--disable-gpu")
349            .arg("--no-sandbox")
350            .arg(&profile_arg)
351            .arg("--hide-scrollbars")
352            .arg("--run-all-compositor-stages-before-draw")
353            .arg("--disable-blink-features=AutomationControlled")
354            .arg("--user-agent=Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/133.0.0.0 Safari/537.36")
355            .arg("--lang=en-US,en")
356            .arg(&window_size_arg)
357            .arg(&screenshot_arg)
358            .arg(&target_arg);
359
360        #[cfg(target_os = "windows")]
361        {
362            use std::os::windows::process::CommandExt;
363            cmd.creation_flags(0x08000000); // CREATE_NO_WINDOW
364        }
365
366        if let Ok(mut child) = cmd.spawn() {
367            let start = std::time::Instant::now();
368            let timeout = std::time::Duration::from_secs(12);
369            loop {
370                if let Ok(Some(_)) = child.try_wait() {
371                    break;
372                }
373                if start.elapsed() > timeout {
374                    let _ = child.kill();
375                    let _ = child.wait();
376                    break;
377                }
378                std::thread::sleep(std::time::Duration::from_millis(50));
379            }
380        }
381
382        if temp_png.exists() {
383            if let Ok(png_bytes) = std::fs::read(&temp_png) {
384                let png_base64 = format!(
385                    "data:image/png;base64,{}",
386                    base64::engine::general_purpose::STANDARD.encode(&png_bytes)
387                );
388                return Some(ScreenshotResult {
389                    width,
390                    height,
391                    svg: String::new(),
392                    layout_wireframe: String::new(),
393                    element_count: 1,
394                    png_bytes,
395                    png_base64,
396                });
397            }
398        }
399        None
400    }
401}
402
403pub struct PageRenderer;
404
405impl PageRenderer {
406    pub async fn render_async(
407        url: &str,
408        title: &str,
409        html_str: &str,
410        interactive: &[InteractiveElement],
411        _search_results: Option<&SearchResults>,
412    ) -> ScreenshotResult {
413        let width = 1280;
414        let height = 720;
415
416        // Try one-shot ultra-light real browser screenshot first
417        if let Some(real_shot) =
418            RealBrowserScreenshot::capture_real_screenshot_async(url, html_str, width, height).await
419        {
420            return real_shot;
421        }
422
423        // Fallback to pure-Rust layout & rasterization engine
424        crate::render::HtmlRenderer::render_html_to_screenshot(
425            url,
426            title,
427            html_str,
428            interactive,
429            width,
430            height,
431        )
432        .await
433        .unwrap_or_else(|_| {
434            Self::render_general_page(url, title, html_str, interactive, width, height)
435        })
436    }
437
438    pub fn render(
439        url: &str,
440        title: &str,
441        html_str: &str,
442        interactive: &[InteractiveElement],
443        _search_results: Option<&SearchResults>,
444    ) -> ScreenshotResult {
445        let width = 1280;
446        let height = 720;
447
448        if let Some(real_shot) =
449            RealBrowserScreenshot::capture_real_screenshot_sync(url, html_str, width, height)
450        {
451            return real_shot;
452        }
453
454        Self::render_general_page(url, title, html_str, interactive, width, height)
455    }
456
457    fn render_general_page(
458        url: &str,
459        title: &str,
460        html_str: &str,
461        interactive: &[InteractiveElement],
462        width: u32,
463        height: u32,
464    ) -> ScreenshotResult {
465        let mut y_offset = 120;
466        let document = Html::parse_document(html_str);
467
468        let mut visual_blocks = Vec::new();
469        if let Ok(sel) = Selector::parse("h1, h2, h3, p, button, a[href], input, li") {
470            for el in document.select(&sel) {
471                let tag = el.value().name();
472                let text = el.text().collect::<Vec<_>>().join(" ").trim().to_string();
473                if text.is_empty() || text.len() < 2 {
474                    continue;
475                }
476
477                let linked_index = interactive.iter().find(|i| i.text == text).map(|i| i.index);
478                visual_blocks.push((tag.to_string(), text, linked_index));
479                if visual_blocks.len() >= 50 {
480                    break;
481                }
482            }
483        }
484
485        let title_escaped = Self::xml_escape(&title.chars().take(30).collect::<String>());
486        let url_escaped = Self::xml_escape(url);
487
488        let mut svg = format!(
489            "<svg xmlns=\"http://www.w3.org/2000/svg\" viewBox=\"0 0 {width} {height}\" width=\"{width}\" height=\"{height}\" style=\"background:#0f172a; font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;\">\n\
490  <rect width=\"{width}\" height=\"80\" fill=\"#1e293b\" />\n\
491  <circle cx=\"25\" cy=\"40\" r=\"7\" fill=\"#ef4444\" />\n\
492  <circle cx=\"45\" cy=\"40\" r=\"7\" fill=\"#f59e0b\" />\n\
493  <circle cx=\"65\" cy=\"40\" r=\"7\" fill=\"#10b981\" />\n\
494  <rect x=\"90\" y=\"24\" width=\"900\" height=\"32\" rx=\"6\" fill=\"#0f172a\" stroke=\"#334155\" stroke-width=\"1\" />\n\
495  <text x=\"105\" y=\"45\" fill=\"#94a3b8\" font-size=\"13\">&#128274; {url_escaped}</text>\n\
496  <text x=\"1010\" y=\"45\" fill=\"#e2e8f0\" font-size=\"13\" font-weight=\"600\">{title_escaped}</text>\n\
497  <g transform=\"translate(60, 100)\">\n"
498        );
499
500        let mut wireframe_lines = Vec::new();
501        wireframe_lines.push(
502            "╔══════════════════════════════════════════════════════════════════════════════╗"
503                .to_string(),
504        );
505        wireframe_lines.push(format!("║ URL: {:<72} ║", url));
506        wireframe_lines.push(format!("║ TITLE: {:<70} ║", title));
507        wireframe_lines.push(
508            "╠══════════════════════════════════════════════════════════════════════════════╣"
509                .to_string(),
510        );
511
512        for (tag, text, index_opt) in &visual_blocks {
513            let truncated_text: String = text.chars().take(80).collect();
514            let escaped_text = Self::xml_escape(&truncated_text);
515
516            match tag.as_str() {
517                "h1" => {
518                    svg.push_str(&format!(
519                        "    <text x=\"0\" y=\"{}\" fill=\"#f8fafc\" font-size=\"24\" font-weight=\"bold\">{}</text>\n",
520                        y_offset, escaped_text
521                    ));
522                    wireframe_lines.push(format!("║ # {:<74} ║", truncated_text));
523                    y_offset += 40;
524                }
525                "h2" => {
526                    svg.push_str(&format!(
527                        "    <text x=\"0\" y=\"{}\" fill=\"#38bdf8\" font-size=\"18\" font-weight=\"bold\">{}</text>\n",
528                        y_offset, escaped_text
529                    ));
530                    wireframe_lines.push(format!("║ ## {:<73} ║", truncated_text));
531                    y_offset += 32;
532                }
533                _ => {
534                    let badge = index_opt.map(|i| format!("[{}] ", i)).unwrap_or_default();
535                    let color = if index_opt.is_some() {
536                        "#60a5fa"
537                    } else {
538                        "#cbd5e1"
539                    };
540                    svg.push_str(&format!(
541                        "    <text x=\"0\" y=\"{}\" fill=\"{}\" font-size=\"13\">{}{}</text>\n",
542                        y_offset, color, badge, escaped_text
543                    ));
544                    if index_opt.is_some() {
545                        wireframe_lines.push(format!("║ [LINK] {}{:<67} ║", badge, truncated_text));
546                    } else {
547                        wireframe_lines.push(format!("║ {:<76} ║", truncated_text));
548                    }
549                    y_offset += 22;
550                }
551            }
552
553            if y_offset > 1100 {
554                break;
555            }
556        }
557
558        svg.push_str("  </g>\n</svg>");
559        wireframe_lines.push(
560            "╚══════════════════════════════════════════════════════════════════════════════╝"
561                .to_string(),
562        );
563
564        let png_bytes = Self::render_png(&svg, width, height).unwrap_or_default();
565        let png_base64 = if !png_bytes.is_empty() {
566            format!(
567                "data:image/png;base64,{}",
568                base64::engine::general_purpose::STANDARD.encode(&png_bytes)
569            )
570        } else {
571            String::new()
572        };
573
574        ScreenshotResult {
575            width,
576            height,
577            svg,
578            layout_wireframe: wireframe_lines.join("\n"),
579            element_count: visual_blocks.len(),
580            png_bytes,
581            png_base64,
582        }
583    }
584
585    pub fn render_png(svg_str: &str, width: u32, height: u32) -> Result<Vec<u8>> {
586        let opt = resvg::usvg::Options {
587            font_family: "sans-serif".to_string(),
588            ..Default::default()
589        };
590        let tree = resvg::usvg::Tree::from_str(svg_str, &opt)?;
591        let mut pixmap = resvg::tiny_skia::Pixmap::new(width, height)
592            .ok_or_else(|| anyhow::anyhow!("Failed to allocate raster pixmap"))?;
593
594        resvg::render(
595            &tree,
596            resvg::tiny_skia::Transform::default(),
597            &mut pixmap.as_mut(),
598        );
599        let png_data = pixmap.encode_png()?;
600        Ok(png_data)
601    }
602
603    fn xml_escape(s: &str) -> String {
604        s.replace('&', "&amp;")
605            .replace('<', "&lt;")
606            .replace('>', "&gt;")
607            .replace('"', "&quot;")
608            .replace('\'', "&apos;")
609    }
610}
headless_engine/dom/screenshot.rs

headless_engine/dom/
screenshot.rs
