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

headless_engine/dom/
screenshot.rs
