Skip to main content

omnivore_cli/git/
output.rs

1use anyhow::{Context, Result};
2use encoding_rs::UTF_8;
3use serde::{Deserialize, Serialize};
4use std::fs;
5use std::io::{self, Write};
6use std::path::{Path, PathBuf};
7
8use super::filter::FilteredFile;
9
10#[derive(Debug, Clone, Copy)]
11pub enum OutputFormat {
12    Json,
13    Text,
14    Directory,
15}
16
17#[derive(Debug, Serialize, Deserialize)]
18pub struct FileContent {
19    pub path: String,
20    pub content: String,
21}
22
23pub struct OutputWriter {
24    format: OutputFormat,
25    #[allow(dead_code)]
26    root_path: PathBuf,
27    output_path: Option<PathBuf>,
28    force_stdout: bool,
29}
30
31impl OutputWriter {
32    pub fn new(format: OutputFormat, root_path: PathBuf) -> Self {
33        Self {
34            format,
35            root_path,
36            output_path: None,
37            force_stdout: false,
38        }
39    }
40
41    pub fn set_output_path(&mut self, path: PathBuf) {
42        self.output_path = Some(path);
43    }
44    
45    pub fn set_stdout_mode(&mut self) {
46        self.force_stdout = true;
47        self.output_path = None;
48    }
49
50    pub async fn write_files(&self, files: Vec<FilteredFile>) -> Result<usize> {
51        match self.format {
52            OutputFormat::Json => self.write_json(files).await,
53            OutputFormat::Text => self.write_text(files).await,
54            OutputFormat::Directory => self.write_directory(files).await,
55        }
56    }
57
58    async fn write_json(&self, files: Vec<FilteredFile>) -> Result<usize> {
59        let mut file_contents = Vec::new();
60        let mut count = 0;
61
62        for file in files {
63            if let Ok(content) = read_file_content(&file.path) {
64                file_contents.push(FileContent {
65                    path: file.relative_path.display().to_string(),
66                    content,
67                });
68                count += 1;
69            }
70        }
71
72        let json = serde_json::to_string_pretty(&file_contents)
73            .context("Failed to serialize to JSON")?;
74
75        if self.force_stdout || self.output_path.is_none() {
76            print!("{}", json);
77            io::stdout().flush()?;
78        } else if let Some(ref output_path) = self.output_path {
79            tokio::fs::write(output_path, json)
80                .await
81                .context("Failed to write JSON to file")?;
82        }
83
84        Ok(count)
85    }
86
87    async fn write_text(&self, files: Vec<FilteredFile>) -> Result<usize> {
88        let mut output = String::new();
89        let mut count = 0;
90
91        for file in files {
92            if let Ok(content) = read_file_content(&file.path) {
93                output.push_str(&format!(
94                    "---\nFile: {}\n---\n{}\n",
95                    file.relative_path.display(),
96                    content
97                ));
98                count += 1;
99            }
100        }
101
102        if self.force_stdout || self.output_path.is_none() {
103            print!("{}", output);
104            io::stdout().flush()?;
105        } else if let Some(ref output_path) = self.output_path {
106            tokio::fs::write(output_path, output)
107                .await
108                .context("Failed to write text to file")?;
109        }
110
111        Ok(count)
112    }
113
114    async fn write_directory(&self, files: Vec<FilteredFile>) -> Result<usize> {
115        let output_dir = self
116            .output_path
117            .as_ref()
118            .ok_or_else(|| anyhow::anyhow!("Output path required for directory format"))?;
119
120        if output_dir.exists() {
121            if !output_dir.is_dir() {
122                return Err(anyhow::anyhow!(
123                    "Output path exists but is not a directory"
124                ));
125            }
126        } else {
127            tokio::fs::create_dir_all(&output_dir)
128                .await
129                .context("Failed to create output directory")?;
130        }
131
132        let mut count = 0;
133        for file in files {
134            let dest_path = output_dir.join(&file.relative_path);
135            
136            if let Some(parent) = dest_path.parent() {
137                tokio::fs::create_dir_all(parent)
138                    .await
139                    .context("Failed to create parent directories")?;
140            }
141
142            match tokio::fs::copy(&file.path, &dest_path).await {
143                Ok(_) => count += 1,
144                Err(e) => {
145                    eprintln!(
146                        "Warning: Failed to copy {}: {}",
147                        file.relative_path.display(),
148                        e
149                    );
150                }
151            }
152        }
153
154        Ok(count)
155    }
156}
157
158fn read_file_content(path: &Path) -> Result<String> {
159    let bytes = fs::read(path).context("Failed to read file")?;
160    
161    let (cow, _, had_errors) = UTF_8.decode(&bytes);
162    
163    if had_errors {
164        return Err(anyhow::anyhow!(
165            "File contains invalid UTF-8: {}",
166            path.display()
167        ));
168    }
169    
170    Ok(cow.into_owned())
171}
172
173#[allow(dead_code)]
174pub fn format_file_size(size: u64) -> String {
175    const UNITS: &[&str] = &["B", "KB", "MB", "GB"];
176    let mut size = size as f64;
177    let mut unit_index = 0;
178
179    while size >= 1024.0 && unit_index < UNITS.len() - 1 {
180        size /= 1024.0;
181        unit_index += 1;
182    }
183
184    if unit_index == 0 {
185        format!("{} {}", size as u64, UNITS[unit_index])
186    } else {
187        format!("{:.2} {}", size, UNITS[unit_index])
188    }
189}