1use std::collections::HashMap;
10
11use ego_tree::NodeRef;
12use scraper::node::Node;
13use scraper::{ElementRef, Html};
14use url::Url;
15
16use crate::extract;
17use crate::types::UrlReference;
18
19pub(crate) struct RefCollector {
22 seen: HashMap<String, usize>,
24 pub(crate) references: Vec<UrlReference>,
25 base: Option<Url>,
26}
27
28impl RefCollector {
29 pub(crate) fn new(base_url: &str) -> Self {
30 Self {
31 seen: HashMap::new(),
32 references: Vec::new(),
33 base: Url::parse(base_url).ok(),
34 }
35 }
36
37 pub(crate) fn resolve(&self, href: &str) -> Option<String> {
42 let href = href.trim();
43 if href.is_empty() || href.starts_with('#') {
44 return None;
45 }
46 let scheme = href.split(':').next().unwrap_or("").to_ascii_lowercase();
47 if matches!(scheme.as_str(), "javascript" | "mailto" | "data" | "tel") {
48 return None;
49 }
50 match &self.base {
51 Some(base) => base.join(href).ok().map(|u| u.to_string()),
52 None => Url::parse(href).ok().map(|u| u.to_string()),
53 }
54 }
55
56 pub(crate) fn index_for(&mut self, url: String, text: &str) -> usize {
58 if let Some(idx) = self.seen.get(&url) {
59 return *idx;
60 }
61 let idx = self.references.len() + 1;
62 self.seen.insert(url.clone(), idx);
63 self.references.push(UrlReference {
64 index: idx,
65 url,
66 text: text.trim().to_string(),
67 });
68 idx
69 }
70}
71
72fn is_block(name: &str) -> bool {
73 matches!(
74 name,
75 "p" | "div"
76 | "section"
77 | "article"
78 | "header"
79 | "footer"
80 | "h1"
81 | "h2"
82 | "h3"
83 | "h4"
84 | "h5"
85 | "h6"
86 | "li"
87 | "ul"
88 | "ol"
89 | "table"
90 | "tr"
91 | "blockquote"
92 | "pre"
93 | "figure"
94 | "aside"
95 | "nav"
96 | "main"
97 )
98}
99
100const CELL_SEPARATOR: &str = " | ";
107
108fn walk(node: NodeRef<Node>, out: &mut String, refs: &mut RefCollector) {
109 match node.value() {
110 Node::Text(t) => out.push_str(&t[..]),
111 Node::Element(el) => {
112 let name = el.name();
113 if super::is_skippable(name) {
114 return;
115 }
116
117 if name == "br" {
118 out.push('\n');
119 return;
120 }
121
122 if name == "a" {
123 let mut inner = String::new();
125 for child in node.children() {
126 walk(child, &mut inner, refs);
127 }
128 let inner = inner.trim().to_string();
129 out.push_str(&inner);
130 if let Some(href) = el.attr("href") {
131 if let Some(resolved) = refs.resolve(href) {
132 let idx = refs.index_for(resolved, &inner);
133 out.push_str(&format!(" [{}]", idx));
134 }
135 }
136 return;
137 }
138
139 if matches!(name, "td" | "th") {
140 if !out.is_empty() && !out.ends_with('\n') {
143 out.push_str(CELL_SEPARATOR);
144 }
145 for child in node.children() {
146 walk(child, out, refs);
147 }
148 return;
149 }
150
151 let block = is_block(name);
152 if block && !out.ends_with('\n') && !out.is_empty() {
153 out.push('\n');
154 }
155 for child in node.children() {
156 walk(child, out, refs);
157 }
158 if block && !out.ends_with('\n') {
159 out.push('\n');
160 }
161 }
162 _ => {}
163 }
164}
165
166pub fn text_with_refs(doc: &Html, base_url: &str) -> (String, Vec<UrlReference>) {
172 let root: ElementRef = match extract::content_root(doc) {
173 Some(el) => el,
174 None => return (String::new(), Vec::new()),
175 };
176
177 let mut refs = RefCollector::new(base_url);
178 let mut out = String::new();
179 for child in root.children() {
180 walk(child, &mut out, &mut refs);
181 }
182 (out, refs.references)
183}
184
185pub fn html_to_text_with_refs(html: &str, base_url: &str) -> (String, Vec<UrlReference>) {
188 text_with_refs(&Html::parse_document(html), base_url)
189}
190
191pub fn render_references(references: &[UrlReference]) -> String {
194 crate::refs::render_block(references)
195}
196
197#[cfg(test)]
198mod tests {
199 use super::*;
200
201 #[test]
202 fn table_cells_are_separated() {
203 let html = "<article><table>\
204 <tr><th>Name</th><th>Type</th></tr>\
205 <tr><td>alpha</td><td>string</td></tr>\
206 </table></article>";
207 let (text, _) = html_to_text_with_refs(html, "https://x.test/");
208 assert!(text.contains("Name | Type"), "text: {text:?}");
209 assert!(text.contains("alpha | string"), "text: {text:?}");
210 }
211
212 #[test]
213 fn unfetchable_schemes_get_no_reference() {
214 let html = r##"<article><p>
215 <a href="javascript:alert(1)">js</a>
216 <a href="mailto:a@b.c">mail</a>
217 <a href="#top">anchor</a>
218 <a href="/ok">ok</a></p></article>"##;
219 let (text, refs) = html_to_text_with_refs(html, "https://x.test/");
220 assert_eq!(refs.len(), 1, "refs: {refs:?}");
221 assert_eq!(refs[0].url, "https://x.test/ok");
222 assert!(text.contains("ok [1]"));
223 }
224}