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