Skip to main content

ctx/
output.rs

1use std::fs;
2use std::io::{self, Write};
3use std::path::Path;
4
5use crate::formatter::{get_formatter, OutputFormat};
6use crate::tree::generate_tree;
7use crate::walker::FileEntry;
8
9/// Result of context generation.
10pub struct ContextResult {
11    pub content: String,
12    pub file_count: usize,
13    pub total_size: u64,
14    /// Byte length of the generated output (used for token estimation).
15    pub output_bytes: usize,
16}
17
18/// Generate context output from file entries.
19pub fn generate_context(
20    root: &Path,
21    entries: &[FileEntry],
22    format: OutputFormat,
23    include_tree: bool,
24    show_sizes: bool,
25) -> io::Result<ContextResult> {
26    let formatter = get_formatter(format);
27
28    // Determine root name for tree
29    let root_name = root
30        .file_name()
31        .map(|s| s.to_string_lossy().to_string())
32        .unwrap_or_else(|| "project".to_string());
33
34    // Generate tree if requested
35    let tree_block = if include_tree {
36        let tree = generate_tree(&root_name, entries, show_sizes);
37        Some(formatter.format_tree(&tree))
38    } else {
39        None
40    };
41
42    // Generate file blocks
43    let mut file_blocks = Vec::new();
44    let mut total_size = 0u64;
45    let mut processed_count = 0usize;
46
47    for entry in entries {
48        match read_file_content(&entry.absolute_path) {
49            Ok(content) => {
50                let block = formatter.format_file(entry, &content);
51                file_blocks.push(block);
52                total_size += entry.size;
53                processed_count += 1;
54            }
55            Err(e) => {
56                eprintln!(
57                    "Warning: could not read {}: {}",
58                    entry.relative_path.display(),
59                    e
60                );
61            }
62        }
63    }
64
65    // Join file blocks
66    let separator = get_separator(format);
67    let files_block = file_blocks.join(&separator);
68
69    // Wrap everything
70    let content = formatter.wrap(tree_block.as_deref(), &files_block);
71
72    Ok(ContextResult {
73        output_bytes: content.len(),
74        content,
75        file_count: processed_count,
76        total_size,
77    })
78}
79
80/// Stream context output, printing each file as it's processed.
81pub fn stream_context(
82    root: &Path,
83    entries: &[FileEntry],
84    format: OutputFormat,
85    include_tree: bool,
86    show_sizes: bool,
87) -> io::Result<ContextResult> {
88    let formatter = get_formatter(format);
89    let mut stdout = io::stdout().lock();
90
91    // Determine root name for tree
92    let root_name = root
93        .file_name()
94        .map(|s| s.to_string_lossy().to_string())
95        .unwrap_or_else(|| "project".to_string());
96
97    // Generate tree if requested
98    let tree_block = if include_tree {
99        let tree = generate_tree(&root_name, entries, show_sizes);
100        Some(formatter.format_tree(&tree))
101    } else {
102        None
103    };
104
105    // Print opening (skip if empty to avoid blank lines in NDJSON)
106    let start = formatter.stream_start(tree_block.as_deref());
107    let mut output_bytes = 0usize;
108    if !start.is_empty() {
109        writeln!(stdout, "{}", start)?;
110        output_bytes += start.len() + 1;
111    }
112
113    // Stream file blocks
114    let mut total_size = 0u64;
115    let mut processed_count = 0usize;
116    let separator = formatter.separator();
117
118    for (i, entry) in entries.iter().enumerate() {
119        match read_file_content(&entry.absolute_path) {
120            Ok(content) => {
121                let block = formatter.format_file(entry, &content);
122                if i > 0 {
123                    write!(stdout, "{}", separator)?;
124                    output_bytes += separator.len();
125                }
126                write!(stdout, "{}", block)?;
127                output_bytes += block.len();
128                stdout.flush()?;
129                total_size += entry.size;
130                processed_count += 1;
131            }
132            Err(e) => {
133                eprintln!(
134                    "Warning: could not read {}: {}",
135                    entry.relative_path.display(),
136                    e
137                );
138            }
139        }
140    }
141
142    // Print closing
143    let end = formatter.stream_end();
144    if !end.is_empty() {
145        writeln!(stdout, "\n{}", end)?;
146        output_bytes += end.len() + 2;
147    } else {
148        writeln!(stdout)?;
149        output_bytes += 1;
150    }
151
152    Ok(ContextResult {
153        content: String::new(), // Not used in streaming mode
154        file_count: processed_count,
155        total_size,
156        output_bytes,
157    })
158}
159
160/// Read file content, handling encoding gracefully.
161fn read_file_content(path: &Path) -> io::Result<String> {
162    let bytes = fs::read(path)?;
163
164    // Try UTF-8 first
165    match String::from_utf8(bytes.clone()) {
166        Ok(s) => Ok(s),
167        Err(_) => {
168            // Fall back to lossy conversion
169            Ok(String::from_utf8_lossy(&bytes).into_owned())
170        }
171    }
172}
173
174/// Get the separator between file blocks based on format.
175fn get_separator(format: OutputFormat) -> String {
176    match format {
177        OutputFormat::Xml => "\n".to_string(),
178        OutputFormat::Markdown => "\n\n".to_string(),
179        OutputFormat::Plain => "\n".to_string(),
180        OutputFormat::Json => ",".to_string(), // Comma for non-streaming JSON array
181    }
182}
183
184#[cfg(test)]
185mod tests {
186    use super::*;
187
188    #[test]
189    fn test_separator_xml() {
190        assert_eq!(get_separator(OutputFormat::Xml), "\n");
191    }
192
193    #[test]
194    fn test_separator_markdown() {
195        assert_eq!(get_separator(OutputFormat::Markdown), "\n\n");
196    }
197}