Skip to main content

every_other_token/
document_parser.rs

1//! Document text extraction and structure detection.
2
3use std::collections::HashMap;
4
5/// The detected format of a document.
6#[derive(Debug, Clone, PartialEq)]
7pub enum DocumentFormat {
8    PlainText,
9    Markdown,
10    Html,
11    Json,
12    Csv,
13    Code(String),
14    Unknown,
15}
16
17/// A section of a parsed document.
18#[derive(Debug, Clone)]
19pub struct DocumentSection {
20    pub title: Option<String>,
21    pub content: String,
22    /// Heading depth: 0 = body text, 1-6 = heading level.
23    pub level: u8,
24    pub line_start: usize,
25    pub line_end: usize,
26}
27
28/// The result of parsing a document.
29#[derive(Debug, Clone)]
30pub struct ParsedDocument {
31    pub format: DocumentFormat,
32    pub sections: Vec<DocumentSection>,
33    pub metadata: HashMap<String, String>,
34    pub word_count: usize,
35    pub char_count: usize,
36    pub detected_language: Option<String>,
37}
38
39/// Stateless document parser.
40pub struct DocumentParser;
41
42impl DocumentParser {
43    /// Heuristically detect the format of the given text.
44    pub fn detect_format(text: &str) -> DocumentFormat {
45        let trimmed = text.trim();
46
47        // JSON: starts with { or [
48        if (trimmed.starts_with('{') && trimmed.ends_with('}'))
49            || (trimmed.starts_with('[') && trimmed.ends_with(']'))
50        {
51            return DocumentFormat::Json;
52        }
53
54        // HTML: contains common HTML tags
55        if trimmed.contains("</") && (trimmed.contains("<html") || trimmed.contains("<div") || trimmed.contains("<p>") || trimmed.contains("<span")) {
56            return DocumentFormat::Html;
57        }
58
59        // Code fence (Markdown code block at top level)
60        if trimmed.starts_with("```") {
61            let lang = trimmed
62                .lines()
63                .next()
64                .unwrap_or("")
65                .trim_start_matches('`')
66                .trim()
67                .to_string();
68            return DocumentFormat::Code(lang);
69        }
70
71        // Markdown: has ATX headers or emphasis
72        let has_md_header = text.lines().any(|l| l.starts_with('#'));
73        let has_md_emphasis = text.contains("**") || text.contains("__");
74        let has_code_fence = text.contains("```");
75        if has_md_header || has_md_emphasis || has_code_fence {
76            return DocumentFormat::Markdown;
77        }
78
79        // CSV: multiple lines with consistent comma counts
80        let lines: Vec<&str> = text.lines().take(5).collect();
81        if lines.len() >= 2 {
82            let counts: Vec<usize> = lines.iter().map(|l| l.matches(',').count()).collect();
83            let first = counts[0];
84            if first > 0 && counts.iter().all(|&c| c == first) {
85                return DocumentFormat::Csv;
86            }
87        }
88
89        DocumentFormat::PlainText
90    }
91
92    /// Parse a Markdown document into sections.
93    pub fn parse_markdown(text: &str) -> ParsedDocument {
94        let mut sections: Vec<DocumentSection> = Vec::new();
95        let mut metadata: HashMap<String, String> = HashMap::new();
96
97        let lines: Vec<&str> = text.lines().collect();
98        let mut i = 0;
99
100        // YAML front matter
101        if lines.first().copied() == Some("---") {
102            i = 1;
103            while i < lines.len() && lines[i] != "---" {
104                let line = lines[i];
105                if let Some(colon) = line.find(':') {
106                    let key = line[..colon].trim().to_string();
107                    let value = line[colon + 1..].trim().to_string();
108                    metadata.insert(key, value);
109                }
110                i += 1;
111            }
112            if i < lines.len() {
113                i += 1; // skip closing ---
114            }
115        }
116
117        let mut current_title: Option<String> = None;
118        let mut current_level: u8 = 0;
119        let mut current_content = String::new();
120        let mut section_start = i;
121
122        while i < lines.len() {
123            let line = lines[i];
124            if let Some(header) = line.strip_prefix('#') {
125                // Count '#' characters
126                let hashes = line.chars().take_while(|&c| c == '#').count() as u8;
127                let title = line.trim_start_matches('#').trim().to_string();
128
129                // Save previous section
130                if !current_content.trim().is_empty() || current_title.is_some() {
131                    sections.push(DocumentSection {
132                        title: current_title.clone(),
133                        content: current_content.trim().to_string(),
134                        level: current_level,
135                        line_start: section_start,
136                        line_end: i,
137                    });
138                }
139
140                current_title = Some(title);
141                current_level = hashes;
142                current_content = String::new();
143                section_start = i;
144                let _ = header; // suppress unused warning
145            } else {
146                current_content.push_str(line);
147                current_content.push('\n');
148            }
149            i += 1;
150        }
151
152        // Final section
153        if !current_content.trim().is_empty() || current_title.is_some() {
154            sections.push(DocumentSection {
155                title: current_title,
156                content: current_content.trim().to_string(),
157                level: current_level,
158                line_start: section_start,
159                line_end: lines.len(),
160            });
161        }
162
163        let wc = Self::word_count(text);
164        let cc = text.chars().count();
165        let lang = Self::detect_language(text);
166
167        ParsedDocument {
168            format: DocumentFormat::Markdown,
169            sections,
170            metadata,
171            word_count: wc,
172            char_count: cc,
173            detected_language: lang,
174        }
175    }
176
177    /// Parse an HTML document into sections.
178    pub fn parse_html(text: &str) -> ParsedDocument {
179        let mut sections: Vec<DocumentSection> = Vec::new();
180        let mut metadata: HashMap<String, String> = HashMap::new();
181
182        // Extract <title>
183        if let Some(start) = text.find("<title>") {
184            if let Some(end) = text.find("</title>") {
185                let title = &text[start + 7..end];
186                metadata.insert("title".to_string(), title.trim().to_string());
187            }
188        }
189
190        // Extract headings h1-h6
191        let plain = Self::strip_html_tags(text);
192        let lines: Vec<&str> = text.lines().collect();
193
194        for (level, tag) in &[
195            (1u8, "h1"),
196            (2, "h2"),
197            (3, "h3"),
198            (4, "h4"),
199            (5, "h5"),
200            (6, "h6"),
201        ] {
202            let open = format!("<{}>", tag);
203            let close = format!("</{}>", tag);
204            let mut search = text;
205            let mut offset = 0;
206            while let Some(start) = search.find(&open) {
207                let abs_start = offset + start;
208                let after_open = abs_start + open.len();
209                if let Some(end_rel) = text[after_open..].find(&close) {
210                    let heading_text = &text[after_open..after_open + end_rel];
211                    let heading_text = Self::strip_html_tags(heading_text);
212                    let line_start = text[..abs_start].matches('\n').count();
213                    sections.push(DocumentSection {
214                        title: Some(heading_text),
215                        content: String::new(),
216                        level: *level,
217                        line_start,
218                        line_end: line_start + 1,
219                    });
220                    offset = after_open + end_rel + close.len();
221                    search = &text[offset..];
222                } else {
223                    break;
224                }
225            }
226        }
227
228        // Sort sections by line_start
229        sections.sort_by_key(|s| s.line_start);
230
231        // Add a body section with plain text
232        if sections.is_empty() {
233            sections.push(DocumentSection {
234                title: None,
235                content: plain.trim().to_string(),
236                level: 0,
237                line_start: 0,
238                line_end: lines.len(),
239            });
240        }
241
242        let wc = Self::word_count(&plain);
243        let cc = plain.chars().count();
244        let lang = Self::detect_language(&plain);
245
246        ParsedDocument {
247            format: DocumentFormat::Html,
248            sections,
249            metadata,
250            word_count: wc,
251            char_count: cc,
252            detected_language: lang,
253        }
254    }
255
256    /// Parse a JSON document, extracting string values and keys as content.
257    pub fn parse_json(text: &str) -> ParsedDocument {
258        let mut content = String::new();
259        let mut in_string = false;
260        let mut escape = false;
261        let mut current_string = String::new();
262
263        for ch in text.chars() {
264            if escape {
265                current_string.push(ch);
266                escape = false;
267                continue;
268            }
269            match ch {
270                '\\' if in_string => escape = true,
271                '"' => {
272                    if in_string {
273                        // end of string
274                        content.push_str(&current_string);
275                        content.push(' ');
276                        current_string.clear();
277                        in_string = false;
278                    } else {
279                        in_string = true;
280                    }
281                }
282                _ if in_string => current_string.push(ch),
283                _ => {}
284            }
285        }
286
287        let wc = Self::word_count(&content);
288        let cc = content.chars().count();
289        let lang = Self::detect_language(&content);
290
291        let section = DocumentSection {
292            title: Some("JSON Content".to_string()),
293            content: content.trim().to_string(),
294            level: 0,
295            line_start: 0,
296            line_end: text.lines().count(),
297        };
298
299        ParsedDocument {
300            format: DocumentFormat::Json,
301            sections: vec![section],
302            metadata: HashMap::new(),
303            word_count: wc,
304            char_count: cc,
305            detected_language: lang,
306        }
307    }
308
309    /// Parse a CSV document, treating the header row as section titles.
310    pub fn parse_csv(text: &str) -> ParsedDocument {
311        let mut sections: Vec<DocumentSection> = Vec::new();
312        let mut lines = text.lines().enumerate();
313
314        let headers: Vec<String> = if let Some((_, header_line)) = lines.next() {
315            header_line.split(',').map(|s| s.trim().to_string()).collect()
316        } else {
317            return ParsedDocument {
318                format: DocumentFormat::Csv,
319                sections: vec![],
320                metadata: HashMap::new(),
321                word_count: 0,
322                char_count: 0,
323                detected_language: None,
324            };
325        };
326
327        let mut rows: Vec<Vec<String>> = Vec::new();
328        for (_, line) in lines {
329            let row: Vec<String> = line.split(',').map(|s| s.trim().to_string()).collect();
330            rows.push(row);
331        }
332
333        for (col_idx, header) in headers.iter().enumerate() {
334            let col_content: Vec<String> = rows
335                .iter()
336                .filter_map(|r| r.get(col_idx).cloned())
337                .filter(|s| !s.is_empty())
338                .collect();
339            let content = col_content.join(", ");
340            sections.push(DocumentSection {
341                title: Some(header.clone()),
342                content,
343                level: 1,
344                line_start: 0,
345                line_end: rows.len() + 1,
346            });
347        }
348
349        let all_content: String = sections.iter().map(|s| s.content.as_str()).collect::<Vec<_>>().join(" ");
350        let wc = Self::word_count(&all_content);
351        let cc = text.chars().count();
352        let lang = Self::detect_language(&all_content);
353
354        ParsedDocument {
355            format: DocumentFormat::Csv,
356            sections,
357            metadata: HashMap::new(),
358            word_count: wc,
359            char_count: cc,
360            detected_language: lang,
361        }
362    }
363
364    /// Auto-detect format and dispatch to the appropriate parser.
365    pub fn parse(text: &str) -> ParsedDocument {
366        match Self::detect_format(text) {
367            DocumentFormat::Markdown => Self::parse_markdown(text),
368            DocumentFormat::Html => Self::parse_html(text),
369            DocumentFormat::Json => Self::parse_json(text),
370            DocumentFormat::Csv => Self::parse_csv(text),
371            _ => {
372                // Plain text: one section
373                let wc = Self::word_count(text);
374                let cc = text.chars().count();
375                let lang = Self::detect_language(text);
376                let lines = text.lines().count();
377                ParsedDocument {
378                    format: DocumentFormat::PlainText,
379                    sections: vec![DocumentSection {
380                        title: None,
381                        content: text.to_string(),
382                        level: 0,
383                        line_start: 0,
384                        line_end: lines,
385                    }],
386                    metadata: HashMap::new(),
387                    word_count: wc,
388                    char_count: cc,
389                    detected_language: lang,
390                }
391            }
392        }
393    }
394
395    /// Extract all URLs (http/https) and Markdown-style links from text.
396    pub fn extract_links(text: &str) -> Vec<String> {
397        let mut links: Vec<String> = Vec::new();
398
399        // Markdown links: [text](url)
400        let mut rest = text;
401        while let Some(open) = rest.find("](") {
402            let after = &rest[open + 2..];
403            if let Some(close) = after.find(')') {
404                let url = &after[..close];
405                if url.starts_with("http://") || url.starts_with("https://") || url.starts_with('/') {
406                    links.push(url.to_string());
407                }
408                rest = &after[close + 1..];
409            } else {
410                break;
411            }
412        }
413
414        // Bare URLs
415        let words = text.split_whitespace();
416        for word in words {
417            let word = word.trim_matches(|c: char| !c.is_alphanumeric() && c != '/' && c != ':' && c != '.' && c != '-' && c != '_' && c != '?' && c != '=' && c != '&' && c != '#' && c != '%');
418            if word.starts_with("http://") || word.starts_with("https://") {
419                if !links.contains(&word.to_string()) {
420                    links.push(word.to_string());
421                }
422            }
423        }
424
425        links
426    }
427
428    /// Extract code blocks as (language, code) pairs.
429    pub fn extract_code_blocks(text: &str) -> Vec<(String, String)> {
430        let mut blocks: Vec<(String, String)> = Vec::new();
431        let mut rest = text;
432
433        while let Some(start) = rest.find("```") {
434            let after_fence = &rest[start + 3..];
435            // Language identifier is on the first line after ```
436            let newline = after_fence.find('\n').unwrap_or(after_fence.len());
437            let lang = after_fence[..newline].trim().to_string();
438            let code_start = newline + 1;
439            if code_start >= after_fence.len() {
440                break;
441            }
442            let code_rest = &after_fence[code_start..];
443            if let Some(end) = code_rest.find("```") {
444                let code = code_rest[..end].to_string();
445                blocks.push((lang, code));
446                rest = &code_rest[end + 3..];
447            } else {
448                break;
449            }
450        }
451
452        blocks
453    }
454
455    /// Simple heuristic language detection based on common words.
456    pub fn detect_language(text: &str) -> Option<String> {
457        let lower = text.to_lowercase();
458        let words: Vec<&str> = lower.split_whitespace().collect();
459
460        let score = |list: &[&str]| -> usize {
461            words.iter().filter(|w| list.contains(w)).count()
462        };
463
464        let english = ["the", "is", "are", "was", "were", "and", "or", "in", "of", "to", "a", "an", "it", "that", "this"];
465        let french = ["le", "la", "les", "est", "sont", "et", "ou", "dans", "de", "du", "un", "une", "il", "elle", "que"];
466        let spanish = ["el", "la", "los", "las", "es", "son", "y", "o", "en", "de", "un", "una", "que", "se", "no"];
467        let german = ["der", "die", "das", "ist", "sind", "und", "oder", "in", "von", "zu", "ein", "eine", "er", "sie", "es"];
468        // Chinese: look for CJK characters
469        let chinese_chars = text.chars().filter(|&c| c >= '\u{4E00}' && c <= '\u{9FFF}').count();
470
471        let scores = [
472            ("English", score(&english)),
473            ("French", score(&french)),
474            ("Spanish", score(&spanish)),
475            ("German", score(&german)),
476        ];
477
478        if chinese_chars > 5 {
479            return Some("Chinese".to_string());
480        }
481
482        let best = scores.iter().max_by_key(|&&(_, s)| s);
483        if let Some(&(lang, s)) = best {
484            if s > 0 {
485                return Some(lang.to_string());
486            }
487        }
488        None
489    }
490
491    /// Count words in text.
492    pub fn word_count(text: &str) -> usize {
493        text.split_whitespace().count()
494    }
495
496    /// Strip HTML tags from a string.
497    pub fn strip_html_tags(html: &str) -> String {
498        let mut result = String::new();
499        let mut in_tag = false;
500        for ch in html.chars() {
501            match ch {
502                '<' => in_tag = true,
503                '>' => {
504                    in_tag = false;
505                    result.push(' ');
506                }
507                _ if !in_tag => result.push(ch),
508                _ => {}
509            }
510        }
511        // Collapse multiple spaces
512        let mut out = String::new();
513        let mut prev_space = false;
514        for ch in result.chars() {
515            if ch.is_whitespace() {
516                if !prev_space {
517                    out.push(' ');
518                }
519                prev_space = true;
520            } else {
521                out.push(ch);
522                prev_space = false;
523            }
524        }
525        out
526    }
527
528    /// Split text into sentences.
529    pub fn split_sentences(text: &str) -> Vec<String> {
530        let mut sentences: Vec<String> = Vec::new();
531        let mut current = String::new();
532
533        for ch in text.chars() {
534            current.push(ch);
535            if ch == '.' || ch == '!' || ch == '?' {
536                let trimmed = current.trim().to_string();
537                if !trimmed.is_empty() {
538                    sentences.push(trimmed);
539                }
540                current.clear();
541            }
542        }
543
544        let trimmed = current.trim().to_string();
545        if !trimmed.is_empty() {
546            sentences.push(trimmed);
547        }
548
549        sentences
550    }
551}
552
553#[cfg(test)]
554mod tests {
555    use super::*;
556
557    #[test]
558    fn test_detect_format_markdown() {
559        let md = "# Hello\n\nSome **bold** text.";
560        assert_eq!(DocumentParser::detect_format(md), DocumentFormat::Markdown);
561    }
562
563    #[test]
564    fn test_detect_format_json() {
565        let json = r#"{"key": "value"}"#;
566        assert_eq!(DocumentParser::detect_format(json), DocumentFormat::Json);
567    }
568
569    #[test]
570    fn test_word_count() {
571        assert_eq!(DocumentParser::word_count("hello world foo"), 3);
572    }
573
574    #[test]
575    fn test_strip_html_tags() {
576        let html = "<h1>Title</h1><p>Body text.</p>";
577        let plain = DocumentParser::strip_html_tags(html);
578        assert!(plain.contains("Title"));
579        assert!(plain.contains("Body text."));
580    }
581
582    #[test]
583    fn test_extract_links() {
584        let text = "Visit [docs](https://example.com) for more info. Also see https://rust-lang.org";
585        let links = DocumentParser::extract_links(text);
586        assert!(!links.is_empty());
587    }
588
589    #[test]
590    fn test_detect_language_english() {
591        let text = "The quick brown fox is and the lazy dog was in the field";
592        assert_eq!(DocumentParser::detect_language(text), Some("English".to_string()));
593    }
594
595    #[test]
596    fn test_parse_markdown() {
597        let md = "# Title\n\nBody text\n\n## Section\n\nMore content.";
598        let doc = DocumentParser::parse_markdown(md);
599        assert_eq!(doc.format, DocumentFormat::Markdown);
600        assert!(!doc.sections.is_empty());
601    }
602
603    #[test]
604    fn test_extract_code_blocks() {
605        let md = "Some text\n```rust\nfn main() {}\n```\nMore text";
606        let blocks = DocumentParser::extract_code_blocks(md);
607        assert_eq!(blocks.len(), 1);
608        assert_eq!(blocks[0].0, "rust");
609    }
610
611    #[test]
612    fn test_split_sentences() {
613        let text = "Hello world. How are you? I am fine!";
614        let sents = DocumentParser::split_sentences(text);
615        assert_eq!(sents.len(), 3);
616    }
617}