context_creator/core/
context_builder.rs

1//! Context creation functionality for LLM consumption
2
3use crate::core::cache::FileCache;
4use crate::core::walker::FileInfo;
5use crate::utils::file_ext::FileType;
6use anyhow::Result;
7use std::collections::HashMap;
8use std::path::Path;
9use std::sync::Arc;
10
11/// Options for generating context for LLM consumption
12#[derive(Debug, Clone)]
13pub struct ContextOptions {
14    /// Maximum tokens allowed in the output
15    pub max_tokens: Option<usize>,
16    /// Include file tree in output
17    pub include_tree: bool,
18    /// Include token count statistics
19    pub include_stats: bool,
20    /// Group files by type
21    pub group_by_type: bool,
22    /// Sort files by priority
23    pub sort_by_priority: bool,
24    /// Template for file headers
25    pub file_header_template: String,
26    /// Template for the document header
27    pub doc_header_template: String,
28    /// Include table of contents
29    pub include_toc: bool,
30    /// Enable enhanced context with file metadata
31    pub enhanced_context: bool,
32}
33
34impl ContextOptions {
35    /// Create ContextOptions from CLI config
36    pub fn from_config(config: &crate::cli::Config) -> Result<Self> {
37        Ok(ContextOptions {
38            max_tokens: config.get_effective_context_tokens(),
39            include_tree: true,
40            include_stats: true,
41            group_by_type: false,
42            sort_by_priority: true,
43            file_header_template: "## {path}".to_string(),
44            doc_header_template: "# Code Context: {directory}".to_string(),
45            include_toc: true,
46            enhanced_context: config.enhanced_context,
47        })
48    }
49}
50
51impl Default for ContextOptions {
52    fn default() -> Self {
53        ContextOptions {
54            max_tokens: None,
55            include_tree: true,
56            include_stats: true,
57            group_by_type: false,
58            sort_by_priority: true,
59            file_header_template: "## {path}".to_string(),
60            doc_header_template: "# Code Context: {directory}".to_string(),
61            include_toc: true,
62            enhanced_context: false,
63        }
64    }
65}
66
67/// Estimate the total size of the markdown output
68fn estimate_output_size(files: &[FileInfo], options: &ContextOptions, cache: &FileCache) -> usize {
69    let mut size = 0;
70
71    // Document header
72    if !options.doc_header_template.is_empty() {
73        size += options.doc_header_template.len() + 50; // Extra for replacements and newlines
74    }
75
76    // Statistics section
77    if options.include_stats {
78        size += 500; // Estimated size for stats
79        size += files.len() * 50; // For file type listing
80    }
81
82    // File tree
83    if options.include_tree {
84        size += 100; // Headers
85        size += files.len() * 100; // Estimated per-file in tree
86    }
87
88    // Table of contents
89    if options.include_toc {
90        size += 50; // Header
91        size += files.len() * 100; // Per-file TOC entry
92    }
93
94    // File contents
95    for file in files {
96        // Header template
97        size +=
98            options.file_header_template.len() + file.relative_path.to_string_lossy().len() + 20;
99
100        // File content + code fence
101        if let Ok(content) = cache.get_or_load(&file.path) {
102            size += content.len() + 20; // Content + fence markers
103        } else {
104            size += file.size as usize; // Fallback to file size
105        }
106    }
107
108    // Add 20% buffer for formatting and unexpected overhead
109    size + (size / 5)
110}
111
112/// Generate markdown from a list of files
113pub fn generate_markdown(
114    files: Vec<FileInfo>,
115    options: ContextOptions,
116    cache: Arc<FileCache>,
117) -> Result<String> {
118    // Pre-allocate string with estimated capacity
119    let estimated_size = estimate_output_size(&files, &options, &cache);
120    let mut output = String::with_capacity(estimated_size);
121
122    // Add document header
123    if !options.doc_header_template.is_empty() {
124        let header = options.doc_header_template.replace("{directory}", ".");
125        output.push_str(&header);
126        output.push_str("\n\n");
127    }
128
129    // Add statistics if requested
130    if options.include_stats {
131        let stats = generate_statistics(&files);
132        output.push_str(&stats);
133        output.push_str("\n\n");
134    }
135
136    // Add file tree if requested
137    if options.include_tree {
138        let tree = generate_file_tree(&files, &options);
139        output.push_str("## File Structure\n\n");
140        output.push_str("```\n");
141        output.push_str(&tree);
142        output.push_str("```\n\n");
143    }
144
145    // Sort files if requested
146    let mut files = files;
147    if options.sort_by_priority {
148        files.sort_by(|a, b| {
149            b.priority
150                .partial_cmp(&a.priority)
151                .unwrap_or(std::cmp::Ordering::Equal)
152                .then_with(|| a.relative_path.cmp(&b.relative_path))
153        });
154    }
155
156    // Add table of contents if requested
157    if options.include_toc {
158        output.push_str("## Table of Contents\n\n");
159        for file in &files {
160            let anchor = path_to_anchor(&file.relative_path);
161            output.push_str(&format!(
162                "- [{path}](#{anchor})\n",
163                path = file.relative_path.display(),
164                anchor = anchor
165            ));
166        }
167        output.push('\n');
168    }
169
170    // Group files if requested
171    if options.group_by_type {
172        let grouped = group_files_by_type(files);
173        for (file_type, group_files) in grouped {
174            output.push_str(&format!("## {} Files\n\n", file_type_display(&file_type)));
175            for file in group_files {
176                append_file_content(&mut output, &file, &options, &cache)?;
177            }
178        }
179    } else {
180        // Add all files
181        for file in files {
182            append_file_content(&mut output, &file, &options, &cache)?;
183        }
184    }
185
186    Ok(output)
187}
188
189/// Append a single file's content to the output
190fn append_file_content(
191    output: &mut String,
192    file: &FileInfo,
193    options: &ContextOptions,
194    cache: &FileCache,
195) -> Result<()> {
196    // Read file content from cache
197    let content = match cache.get_or_load(&file.path) {
198        Ok(content) => content,
199        Err(e) => {
200            eprintln!(
201                "Warning: Could not read file {}: {}",
202                file.path.display(),
203                e
204            );
205            return Ok(());
206        }
207    };
208
209    // Add file header with optional metadata
210    let path_with_metadata = if options.enhanced_context {
211        format!(
212            "{} ({}, {})",
213            file.relative_path.display(),
214            format_size(file.size),
215            file_type_display(&file.file_type)
216        )
217    } else {
218        file.relative_path.display().to_string()
219    };
220
221    let header = options
222        .file_header_template
223        .replace("{path}", &path_with_metadata);
224    output.push_str(&header);
225    output.push_str("\n\n");
226
227    // Add semantic information if available
228    if !file.imports.is_empty() {
229        output.push_str("Imports: ");
230        let import_names: Vec<String> = file
231            .imports
232            .iter()
233            .map(|p| {
234                let filename = p.file_name().and_then(|n| n.to_str()).unwrap_or("");
235
236                // For Python __init__.py files, use the parent directory name
237                if filename == "__init__.py" {
238                    p.parent()
239                        .and_then(|parent| parent.file_name())
240                        .and_then(|n| n.to_str())
241                        .unwrap_or("unknown")
242                        .to_string()
243                } else {
244                    // Remove common extensions
245                    filename
246                        .strip_suffix(".py")
247                        .or_else(|| filename.strip_suffix(".rs"))
248                        .or_else(|| filename.strip_suffix(".js"))
249                        .or_else(|| filename.strip_suffix(".ts"))
250                        .unwrap_or(filename)
251                        .to_string()
252                }
253            })
254            .collect();
255        output.push_str(&format!("{}\n\n", import_names.join(", ")));
256    }
257
258    if !file.imported_by.is_empty() {
259        output.push_str("Imported by: ");
260        let imported_by_names: Vec<String> = file
261            .imported_by
262            .iter()
263            .map(|p| {
264                p.file_name()
265                    .and_then(|n| n.to_str())
266                    .unwrap_or_else(|| p.to_str().unwrap_or("unknown"))
267                    .to_string()
268            })
269            .collect();
270        output.push_str(&format!("{}\n\n", imported_by_names.join(", ")));
271    }
272
273    if !file.function_calls.is_empty() {
274        output.push_str("Function calls: ");
275        let function_names: Vec<String> = file
276            .function_calls
277            .iter()
278            .map(|fc| {
279                if let Some(module) = &fc.module {
280                    format!("{}.{}", module, fc.name)
281                } else {
282                    fc.name.clone()
283                }
284            })
285            .collect();
286        output.push_str(&format!("{}\n\n", function_names.join(", ")));
287    }
288
289    if !file.type_references.is_empty() {
290        output.push_str("Type references: ");
291        let type_names: Vec<String> = file
292            .type_references
293            .iter()
294            .map(|tr| {
295                if let Some(module) = &tr.module {
296                    format!("{}.{}", module, tr.name)
297                } else {
298                    tr.name.clone()
299                }
300            })
301            .collect();
302        output.push_str(&format!("{}\n\n", type_names.join(", ")));
303    }
304
305    // Add language hint for syntax highlighting
306    let language = get_language_hint(&file.file_type);
307    output.push_str(&format!("```{language}\n"));
308    output.push_str(&content);
309    if !content.ends_with('\n') {
310        output.push('\n');
311    }
312    output.push_str("```\n\n");
313
314    Ok(())
315}
316
317/// Generate statistics about the files
318fn generate_statistics(files: &[FileInfo]) -> String {
319    let total_files = files.len();
320    let total_size: u64 = files.iter().map(|f| f.size).sum();
321
322    // Count by file type
323    let mut type_counts: HashMap<FileType, usize> = HashMap::new();
324    for file in files {
325        *type_counts.entry(file.file_type.clone()).or_insert(0) += 1;
326    }
327
328    // Pre-allocate with estimated capacity
329    let mut stats = String::with_capacity(500 + type_counts.len() * 50);
330    stats.push_str("## Statistics\n\n");
331    stats.push_str(&format!("- Total files: {total_files}\n"));
332    stats.push_str(&format!(
333        "- Total size: {} bytes\n",
334        format_size(total_size)
335    ));
336    stats.push_str("\n### Files by type:\n");
337
338    let mut types: Vec<_> = type_counts.into_iter().collect();
339    types.sort_by_key(|(_, count)| std::cmp::Reverse(*count));
340
341    for (file_type, count) in types {
342        stats.push_str(&format!("- {}: {}\n", file_type_display(&file_type), count));
343    }
344
345    stats
346}
347
348/// Generate a file tree representation
349fn generate_file_tree(files: &[FileInfo], options: &ContextOptions) -> String {
350    use std::collections::{BTreeMap, HashMap};
351
352    #[derive(Default)]
353    struct TreeNode {
354        files: Vec<String>,
355        dirs: BTreeMap<String, TreeNode>,
356    }
357
358    let mut root = TreeNode::default();
359
360    // Create a lookup map from relative path to FileInfo for metadata
361    let file_lookup: HashMap<String, &FileInfo> = files
362        .iter()
363        .map(|f| (f.relative_path.to_string_lossy().to_string(), f))
364        .collect();
365
366    // Build tree structure
367    for file in files {
368        let parts: Vec<_> = file
369            .relative_path
370            .components()
371            .map(|c| c.as_os_str().to_string_lossy().to_string())
372            .collect();
373
374        let mut current = &mut root;
375        for (i, part) in parts.iter().enumerate() {
376            if i == parts.len() - 1 {
377                // File
378                current.files.push(part.clone());
379            } else {
380                // Directory
381                current = current.dirs.entry(part.clone()).or_default();
382            }
383        }
384    }
385
386    // Render tree
387    fn render_tree(
388        node: &TreeNode,
389        prefix: &str,
390        _is_last: bool,
391        current_path: &str,
392        file_lookup: &HashMap<String, &FileInfo>,
393        options: &ContextOptions,
394    ) -> String {
395        // Pre-allocate with estimated size
396        let estimated_size = (node.dirs.len() + node.files.len()) * 100;
397        let mut output = String::with_capacity(estimated_size);
398
399        // Render directories
400        let dir_count = node.dirs.len();
401        for (i, (name, child)) in node.dirs.iter().enumerate() {
402            let is_last_dir = i == dir_count - 1 && node.files.is_empty();
403            let connector = if is_last_dir {
404                "└── "
405            } else {
406                "├── "
407            };
408            let extension = if is_last_dir { "    " } else { "│   " };
409
410            output.push_str(&format!("{prefix}{connector}{name}/\n"));
411            let child_path = if current_path.is_empty() {
412                name.clone()
413            } else {
414                format!("{current_path}/{name}")
415            };
416            output.push_str(&render_tree(
417                child,
418                &format!("{prefix}{extension}"),
419                is_last_dir,
420                &child_path,
421                file_lookup,
422                options,
423            ));
424        }
425
426        // Render files
427        let file_count = node.files.len();
428        for (i, name) in node.files.iter().enumerate() {
429            let is_last_file = i == file_count - 1;
430            let connector = if is_last_file {
431                "└── "
432            } else {
433                "├── "
434            };
435
436            let file_path = if current_path.is_empty() {
437                name.clone()
438            } else {
439                format!("{current_path}/{name}")
440            };
441
442            // Include metadata if enhanced context is enabled
443            let display_name = if options.enhanced_context {
444                if let Some(file_info) = file_lookup.get(&file_path) {
445                    format!(
446                        "{} ({}, {})",
447                        name,
448                        format_size(file_info.size),
449                        file_type_display(&file_info.file_type)
450                    )
451                } else {
452                    name.clone()
453                }
454            } else {
455                name.clone()
456            };
457
458            output.push_str(&format!("{prefix}{connector}{display_name}\n"));
459        }
460
461        output
462    }
463
464    // Pre-allocate output string
465    let mut output = String::with_capacity(files.len() * 100 + 10);
466    output.push_str(".\n");
467    output.push_str(&render_tree(&root, "", true, "", &file_lookup, options));
468    output
469}
470
471/// Group files by their type
472fn group_files_by_type(files: Vec<FileInfo>) -> Vec<(FileType, Vec<FileInfo>)> {
473    let mut groups: HashMap<FileType, Vec<FileInfo>> = HashMap::new();
474
475    for file in files {
476        groups.entry(file.file_type.clone()).or_default().push(file);
477    }
478
479    let mut result: Vec<_> = groups.into_iter().collect();
480    result.sort_by_key(|(file_type, _)| file_type_priority(file_type));
481    result
482}
483
484/// Get display name for file type
485fn file_type_display(file_type: &FileType) -> &'static str {
486    match file_type {
487        FileType::Rust => "Rust",
488        FileType::Python => "Python",
489        FileType::JavaScript => "JavaScript",
490        FileType::TypeScript => "TypeScript",
491        FileType::Go => "Go",
492        FileType::Java => "Java",
493        FileType::Cpp => "C++",
494        FileType::C => "C",
495        FileType::CSharp => "C#",
496        FileType::Ruby => "Ruby",
497        FileType::Php => "PHP",
498        FileType::Swift => "Swift",
499        FileType::Kotlin => "Kotlin",
500        FileType::Scala => "Scala",
501        FileType::Haskell => "Haskell",
502        FileType::Dart => "Dart",
503        FileType::Lua => "Lua",
504        FileType::R => "R",
505        FileType::Julia => "Julia",
506        FileType::Elixir => "Elixir",
507        FileType::Elm => "Elm",
508        FileType::Markdown => "Markdown",
509        FileType::Json => "JSON",
510        FileType::Yaml => "YAML",
511        FileType::Toml => "TOML",
512        FileType::Xml => "XML",
513        FileType::Html => "HTML",
514        FileType::Css => "CSS",
515        FileType::Text => "Text",
516        FileType::Other => "Other",
517    }
518}
519
520/// Get language hint for syntax highlighting
521fn get_language_hint(file_type: &FileType) -> &'static str {
522    match file_type {
523        FileType::Rust => "rust",
524        FileType::Python => "python",
525        FileType::JavaScript => "javascript",
526        FileType::TypeScript => "typescript",
527        FileType::Go => "go",
528        FileType::Java => "java",
529        FileType::Cpp => "cpp",
530        FileType::C => "c",
531        FileType::CSharp => "csharp",
532        FileType::Ruby => "ruby",
533        FileType::Php => "php",
534        FileType::Swift => "swift",
535        FileType::Kotlin => "kotlin",
536        FileType::Scala => "scala",
537        FileType::Haskell => "haskell",
538        FileType::Dart => "dart",
539        FileType::Lua => "lua",
540        FileType::R => "r",
541        FileType::Julia => "julia",
542        FileType::Elixir => "elixir",
543        FileType::Elm => "elm",
544        FileType::Markdown => "markdown",
545        FileType::Json => "json",
546        FileType::Yaml => "yaml",
547        FileType::Toml => "toml",
548        FileType::Xml => "xml",
549        FileType::Html => "html",
550        FileType::Css => "css",
551        FileType::Text => "text",
552        FileType::Other => "",
553    }
554}
555
556/// Get priority for file type ordering
557fn file_type_priority(file_type: &FileType) -> u8 {
558    match file_type {
559        FileType::Rust => 1,
560        FileType::Python => 2,
561        FileType::JavaScript => 3,
562        FileType::TypeScript => 3,
563        FileType::Go => 4,
564        FileType::Java => 5,
565        FileType::Cpp => 6,
566        FileType::C => 7,
567        FileType::CSharp => 8,
568        FileType::Ruby => 9,
569        FileType::Php => 10,
570        FileType::Swift => 11,
571        FileType::Kotlin => 12,
572        FileType::Scala => 13,
573        FileType::Haskell => 14,
574        FileType::Dart => 15,
575        FileType::Lua => 16,
576        FileType::R => 17,
577        FileType::Julia => 18,
578        FileType::Elixir => 19,
579        FileType::Elm => 20,
580        FileType::Markdown => 21,
581        FileType::Json => 22,
582        FileType::Yaml => 23,
583        FileType::Toml => 24,
584        FileType::Xml => 25,
585        FileType::Html => 26,
586        FileType::Css => 27,
587        FileType::Text => 28,
588        FileType::Other => 29,
589    }
590}
591
592/// Convert path to anchor-friendly string
593fn path_to_anchor(path: &Path) -> String {
594    path.display()
595        .to_string()
596        .replace(['/', '\\', '.', ' '], "-")
597        .to_lowercase()
598}
599
600/// Format file size in human-readable format
601fn format_size(size: u64) -> String {
602    const UNITS: &[&str] = &["B", "KB", "MB", "GB"];
603    let mut size = size as f64;
604    let mut unit_index = 0;
605
606    while size >= 1024.0 && unit_index < UNITS.len() - 1 {
607        size /= 1024.0;
608        unit_index += 1;
609    }
610
611    if unit_index == 0 {
612        format!("{} {}", size as u64, UNITS[unit_index])
613    } else {
614        format!("{:.2} {}", size, UNITS[unit_index])
615    }
616}
617
618#[cfg(test)]
619mod tests {
620    use super::*;
621    use std::path::PathBuf;
622
623    fn create_test_cache() -> Arc<FileCache> {
624        Arc::new(FileCache::new())
625    }
626
627    #[test]
628    fn test_format_size() {
629        assert_eq!(format_size(512), "512 B");
630        assert_eq!(format_size(1024), "1.00 KB");
631        assert_eq!(format_size(1536), "1.50 KB");
632        assert_eq!(format_size(1048576), "1.00 MB");
633    }
634
635    #[test]
636    fn test_path_to_anchor() {
637        assert_eq!(path_to_anchor(Path::new("src/main.rs")), "src-main-rs");
638        assert_eq!(path_to_anchor(Path::new("test file.txt")), "test-file-txt");
639    }
640
641    #[test]
642    fn test_file_type_display() {
643        assert_eq!(file_type_display(&FileType::Rust), "Rust");
644        assert_eq!(file_type_display(&FileType::Python), "Python");
645    }
646
647    #[test]
648    fn test_generate_statistics() {
649        let files = vec![
650            FileInfo {
651                path: PathBuf::from("test1.rs"),
652                relative_path: PathBuf::from("test1.rs"),
653                size: 100,
654                file_type: FileType::Rust,
655                priority: 1.0,
656                imports: Vec::new(),
657                imported_by: Vec::new(),
658                function_calls: Vec::new(),
659                type_references: Vec::new(),
660            },
661            FileInfo {
662                path: PathBuf::from("test2.py"),
663                relative_path: PathBuf::from("test2.py"),
664                size: 200,
665                file_type: FileType::Python,
666                priority: 0.9,
667                imports: Vec::new(),
668                imported_by: Vec::new(),
669                function_calls: Vec::new(),
670                type_references: Vec::new(),
671            },
672        ];
673
674        let stats = generate_statistics(&files);
675        assert!(stats.contains("Total files: 2"));
676        assert!(stats.contains("Total size: 300 B"));
677        assert!(stats.contains("Rust: 1"));
678        assert!(stats.contains("Python: 1"));
679    }
680
681    #[test]
682    fn test_generate_statistics_empty() {
683        let files = vec![];
684        let stats = generate_statistics(&files);
685        assert!(stats.contains("Total files: 0"));
686        assert!(stats.contains("Total size: 0 B"));
687    }
688
689    #[test]
690    fn test_generate_statistics_large_files() {
691        let files = vec![
692            FileInfo {
693                path: PathBuf::from("large.rs"),
694                relative_path: PathBuf::from("large.rs"),
695                size: 2_000_000, // 2MB
696                file_type: FileType::Rust,
697                priority: 1.0,
698                imports: Vec::new(),
699                imported_by: Vec::new(),
700                function_calls: Vec::new(),
701                type_references: Vec::new(),
702            },
703            FileInfo {
704                path: PathBuf::from("huge.py"),
705                relative_path: PathBuf::from("huge.py"),
706                size: 50_000_000, // 50MB
707                file_type: FileType::Python,
708                priority: 0.9,
709                imports: Vec::new(),
710                imported_by: Vec::new(),
711                function_calls: Vec::new(),
712                type_references: Vec::new(),
713            },
714        ];
715
716        let stats = generate_statistics(&files);
717        assert!(stats.contains("Total files: 2"));
718        assert!(stats.contains("MB bytes")); // Just check that it's in MB
719        assert!(stats.contains("Python: 1"));
720        assert!(stats.contains("Rust: 1"));
721    }
722
723    #[test]
724    fn test_generate_file_tree_with_grouping() {
725        let files = vec![
726            FileInfo {
727                path: PathBuf::from("src/main.rs"),
728                relative_path: PathBuf::from("src/main.rs"),
729                size: 1000,
730                file_type: FileType::Rust,
731                priority: 1.5,
732                imports: Vec::new(),
733                imported_by: Vec::new(),
734                function_calls: Vec::new(),
735                type_references: Vec::new(),
736            },
737            FileInfo {
738                path: PathBuf::from("src/lib.rs"),
739                relative_path: PathBuf::from("src/lib.rs"),
740                size: 2000,
741                file_type: FileType::Rust,
742                priority: 1.2,
743                imports: Vec::new(),
744                imported_by: Vec::new(),
745                function_calls: Vec::new(),
746                type_references: Vec::new(),
747            },
748            FileInfo {
749                path: PathBuf::from("tests/test.rs"),
750                relative_path: PathBuf::from("tests/test.rs"),
751                size: 500,
752                file_type: FileType::Rust,
753                priority: 0.8,
754                imports: Vec::new(),
755                imported_by: Vec::new(),
756                function_calls: Vec::new(),
757                type_references: Vec::new(),
758            },
759        ];
760
761        let options = ContextOptions::default();
762        let tree = generate_file_tree(&files, &options);
763        assert!(tree.contains("src/"));
764        assert!(tree.contains("tests/"));
765        assert!(tree.contains("main.rs"));
766        assert!(tree.contains("lib.rs"));
767        assert!(tree.contains("test.rs"));
768    }
769
770    #[test]
771    fn test_context_options_from_config() {
772        use crate::cli::Config;
773        use tempfile::TempDir;
774
775        let temp_dir = TempDir::new().unwrap();
776        let config = Config {
777            paths: Some(vec![temp_dir.path().to_path_buf()]),
778            max_tokens: Some(100000),
779            ..Config::default()
780        };
781
782        let options = ContextOptions::from_config(&config).unwrap();
783        assert_eq!(options.max_tokens, Some(100000));
784        assert!(options.include_tree);
785        assert!(options.include_stats);
786        assert!(!options.group_by_type); // Default is false according to implementation
787    }
788
789    #[test]
790    fn test_generate_markdown_structure_headers() {
791        let files = vec![];
792
793        let options = ContextOptions {
794            max_tokens: None,
795            include_tree: true,
796            include_stats: true,
797            group_by_type: true,
798            sort_by_priority: true,
799            file_header_template: "## {path}".to_string(),
800            doc_header_template: "# Code Context".to_string(),
801            include_toc: true,
802            enhanced_context: false,
803        };
804
805        let cache = create_test_cache();
806        let markdown = generate_markdown(files, options, cache).unwrap();
807
808        // Check that main structure is present even with no files
809        assert!(markdown.contains("# Code Context"));
810        assert!(markdown.contains("## Statistics"));
811    }
812
813    #[test]
814    fn test_enhanced_tree_generation_with_metadata() {
815        use crate::core::walker::FileInfo;
816        use crate::utils::file_ext::FileType;
817        use std::path::PathBuf;
818
819        let files = vec![
820            FileInfo {
821                path: PathBuf::from("src/main.rs"),
822                relative_path: PathBuf::from("src/main.rs"),
823                size: 145,
824                file_type: FileType::Rust,
825                priority: 1.5,
826                imports: Vec::new(),
827                imported_by: Vec::new(),
828                function_calls: Vec::new(),
829                type_references: Vec::new(),
830            },
831            FileInfo {
832                path: PathBuf::from("src/lib.rs"),
833                relative_path: PathBuf::from("src/lib.rs"),
834                size: 89,
835                file_type: FileType::Rust,
836                priority: 1.2,
837                imports: Vec::new(),
838                imported_by: Vec::new(),
839                function_calls: Vec::new(),
840                type_references: Vec::new(),
841            },
842        ];
843
844        let options = ContextOptions {
845            max_tokens: None,
846            include_tree: true,
847            include_stats: true,
848            group_by_type: false,
849            sort_by_priority: true,
850            file_header_template: "## {path}".to_string(),
851            doc_header_template: "# Code Context".to_string(),
852            include_toc: true,
853            enhanced_context: true,
854        };
855
856        let cache = create_test_cache();
857        let markdown = generate_markdown(files, options, cache).unwrap();
858
859        // Should include file sizes and types in tree
860        assert!(markdown.contains("main.rs (145 B, Rust)"));
861        assert!(markdown.contains("lib.rs (89 B, Rust)"));
862    }
863
864    #[test]
865    fn test_enhanced_file_headers_with_metadata() {
866        use crate::core::walker::FileInfo;
867        use crate::utils::file_ext::FileType;
868        use std::path::PathBuf;
869
870        let files = vec![FileInfo {
871            path: PathBuf::from("src/main.rs"),
872            relative_path: PathBuf::from("src/main.rs"),
873            size: 145,
874            file_type: FileType::Rust,
875            priority: 1.5,
876            imports: Vec::new(),
877            imported_by: Vec::new(),
878            function_calls: Vec::new(),
879            type_references: Vec::new(),
880        }];
881
882        let options = ContextOptions {
883            max_tokens: None,
884            include_tree: true,
885            include_stats: true,
886            group_by_type: false,
887            sort_by_priority: true,
888            file_header_template: "## {path}".to_string(),
889            doc_header_template: "# Code Context".to_string(),
890            include_toc: true,
891            enhanced_context: true,
892        };
893
894        let cache = create_test_cache();
895        let markdown = generate_markdown(files, options, cache).unwrap();
896
897        // Should include metadata in file headers
898        assert!(markdown.contains("## src/main.rs (145 B, Rust)"));
899    }
900
901    #[test]
902    fn test_basic_mode_unchanged() {
903        use crate::core::walker::FileInfo;
904        use crate::utils::file_ext::FileType;
905        use std::path::PathBuf;
906
907        let files = vec![FileInfo {
908            path: PathBuf::from("src/main.rs"),
909            relative_path: PathBuf::from("src/main.rs"),
910            size: 145,
911            file_type: FileType::Rust,
912            priority: 1.5,
913            imports: Vec::new(),
914            imported_by: Vec::new(),
915            function_calls: Vec::new(),
916            type_references: Vec::new(),
917        }];
918
919        let options = ContextOptions {
920            max_tokens: None,
921            include_tree: true,
922            include_stats: true,
923            group_by_type: false,
924            sort_by_priority: true,
925            file_header_template: "## {path}".to_string(),
926            doc_header_template: "# Code Context".to_string(),
927            include_toc: true,
928            enhanced_context: false,
929        };
930
931        let cache = create_test_cache();
932        let markdown = generate_markdown(files, options, cache).unwrap();
933
934        // Should NOT include metadata - backward compatibility
935        assert!(markdown.contains("## src/main.rs"));
936        assert!(!markdown.contains("## src/main.rs (145 B, Rust)"));
937        assert!(markdown.contains("main.rs") && !markdown.contains("main.rs (145 B, Rust)"));
938    }
939}
context_creator/core/context_builder.rs

context_creator/core/
context_builder.rs
