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

headless_engine/dom/
screenshot.rs
