1use anyhow::Result;
2use std::collections::HashMap;
3use std::path::Path;
4
5use super::detector::CodebaseInfo;
6use super::filter::FilteredFile;
7
8pub struct CodeOrganizer {
9 codebase_info: CodebaseInfo,
10 files: Vec<FilteredFile>,
11}
12
13impl CodeOrganizer {
14 pub fn new(codebase_info: CodebaseInfo, files: Vec<FilteredFile>) -> Self {
15 Self {
16 codebase_info,
17 files,
18 }
19 }
20
21 pub fn organize(&self) -> OrganizedCode {
22 let mut organized = OrganizedCode {
23 metadata: self.generate_metadata(),
24 sections: Vec::new(),
25 };
26
27 let categorized = self.categorize_files();
28
29 for (category, files) in categorized {
30 if !files.is_empty() {
31 organized.sections.push(CodeSection {
32 name: category.clone(),
33 description: self.get_category_description(&category),
34 files: files.iter().cloned().cloned().collect(),
35 });
36 }
37 }
38
39 organized.sections.sort_by(|a, b| {
40 let order = self.get_section_priority(&a.name);
41 let order_b = self.get_section_priority(&b.name);
42 order.cmp(&order_b)
43 });
44
45 organized
46 }
47
48 fn generate_metadata(&self) -> ProjectMetadata {
49 ProjectMetadata {
50 project_type: format!("{:?}", self.codebase_info.project_type),
51 description: self.codebase_info.description.clone(),
52 main_language: self
53 .codebase_info
54 .main_language
55 .as_ref()
56 .map(|l| format!("{:?}", l))
57 .unwrap_or_else(|| "Unknown".to_string()),
58 frameworks: self
59 .codebase_info
60 .frameworks
61 .iter()
62 .map(|f| format!("{:?}", f))
63 .collect(),
64 build_tools: self
65 .codebase_info
66 .build_tools
67 .iter()
68 .map(|b| format!("{:?}", b))
69 .collect(),
70 total_files: self.files.len(),
71 }
72 }
73
74 fn categorize_files(&self) -> HashMap<String, Vec<&FilteredFile>> {
75 let mut categories: HashMap<String, Vec<&FilteredFile>> = HashMap::new();
76
77 for file in &self.files {
78 let category = self.determine_category(&file.relative_path);
79 categories
80 .entry(category)
81 .or_insert_with(Vec::new)
82 .push(file);
83 }
84
85 for files in categories.values_mut() {
86 files.sort_by(|a, b| a.relative_path.cmp(&b.relative_path));
87 }
88
89 categories
90 }
91
92 fn determine_category(&self, path: &Path) -> String {
93 let path_str = path.to_string_lossy().to_lowercase();
94 let file_name = path
95 .file_name()
96 .and_then(|n| n.to_str())
97 .unwrap_or("")
98 .to_lowercase();
99
100 if path_str.contains("test") || path_str.contains("spec") {
101 return "Tests".to_string();
102 }
103
104 if file_name == "readme.md"
105 || file_name == "license"
106 || file_name == "contributing.md"
107 || file_name == "changelog.md"
108 {
109 return "Documentation".to_string();
110 }
111
112 if file_name == "dockerfile"
113 || file_name == "docker-compose.yml"
114 || file_name == ".dockerignore"
115 || path_str.contains("k8s/")
116 || path_str.contains("kubernetes/")
117 {
118 return "Infrastructure".to_string();
119 }
120
121 if file_name.starts_with('.')
122 || file_name == "package.json"
123 || file_name == "cargo.toml"
124 || file_name == "pyproject.toml"
125 || file_name == "go.mod"
126 || file_name == "pom.xml"
127 || file_name == "build.gradle"
128 || file_name == "gemfile"
129 || file_name == "composer.json"
130 || file_name == "makefile"
131 || file_name == "cmakelists.txt"
132 {
133 return "Configuration".to_string();
134 }
135
136 if path_str.contains("migrations/") || path_str.contains("schema") {
137 return "Database".to_string();
138 }
139
140 if let Some(ext) = path.extension().and_then(|e| e.to_str()) {
141 match ext {
142 "html" | "htm" => return "Templates".to_string(),
143 "css" | "scss" | "sass" | "less" => return "Styles".to_string(),
144 "sql" => return "Database".to_string(),
145 "yml" | "yaml" if !file_name.contains("docker") => {
146 return "Configuration".to_string()
147 }
148 _ => {}
149 }
150 }
151
152 let components = path.components().collect::<Vec<_>>();
153 if components.len() > 1 {
154 if let Some(first_dir) = components.first() {
155 let dir_str = first_dir.as_os_str().to_string_lossy().to_lowercase();
156 return match dir_str.as_str() {
157 "src" | "lib" => "Source Code".to_string(),
158 "app" => "Application".to_string(),
159 "pages" => "Pages".to_string(),
160 "components" => "Components".to_string(),
161 "utils" | "helpers" => "Utilities".to_string(),
162 "services" => "Services".to_string(),
163 "api" => "API".to_string(),
164 "models" => "Models".to_string(),
165 "controllers" => "Controllers".to_string(),
166 "views" => "Views".to_string(),
167 "public" | "static" => "Static Assets".to_string(),
168 "scripts" => "Scripts".to_string(),
169 "bin" | "cmd" => "Binaries".to_string(),
170 _ => "Source Code".to_string(),
171 };
172 }
173 }
174
175 "Source Code".to_string()
176 }
177
178 fn get_category_description(&self, category: &str) -> String {
179 match category {
180 "Configuration" => "Project configuration and build files".to_string(),
181 "Source Code" => "Main application source code".to_string(),
182 "Tests" => "Test files and specifications".to_string(),
183 "Documentation" => "Project documentation and guides".to_string(),
184 "Infrastructure" => "Deployment and infrastructure configuration".to_string(),
185 "Database" => "Database schemas and migrations".to_string(),
186 "Templates" => "HTML templates and views".to_string(),
187 "Styles" => "CSS and styling files".to_string(),
188 "Components" => "Reusable UI components".to_string(),
189 "Pages" => "Application pages and routes".to_string(),
190 "API" => "API endpoints and handlers".to_string(),
191 "Services" => "Business logic and service layers".to_string(),
192 "Models" => "Data models and entities".to_string(),
193 "Controllers" => "Request controllers and handlers".to_string(),
194 "Views" => "View templates and presentations".to_string(),
195 "Utilities" => "Helper functions and utilities".to_string(),
196 "Static Assets" => "Static files and resources".to_string(),
197 "Scripts" => "Build and utility scripts".to_string(),
198 "Binaries" => "Executable files and commands".to_string(),
199 "Application" => "Application entry points and core logic".to_string(),
200 _ => format!("{} files", category),
201 }
202 }
203
204 fn get_section_priority(&self, section: &str) -> usize {
205 match section {
206 "Documentation" => 0,
207 "Configuration" => 1,
208 "Source Code" => 2,
209 "Application" => 3,
210 "Pages" => 4,
211 "Components" => 5,
212 "API" => 6,
213 "Services" => 7,
214 "Models" => 8,
215 "Controllers" => 9,
216 "Views" => 10,
217 "Utilities" => 11,
218 "Templates" => 12,
219 "Styles" => 13,
220 "Database" => 14,
221 "Tests" => 15,
222 "Scripts" => 16,
223 "Infrastructure" => 17,
224 "Static Assets" => 18,
225 "Binaries" => 19,
226 _ => 99,
227 }
228 }
229}
230
231#[derive(Debug)]
232pub struct OrganizedCode {
233 pub metadata: ProjectMetadata,
234 pub sections: Vec<CodeSection>,
235}
236
237#[derive(Debug)]
238pub struct ProjectMetadata {
239 pub project_type: String,
240 pub description: String,
241 pub main_language: String,
242 pub frameworks: Vec<String>,
243 pub build_tools: Vec<String>,
244 pub total_files: usize,
245}
246
247#[derive(Debug)]
248pub struct CodeSection {
249 pub name: String,
250 pub description: String,
251 pub files: Vec<FilteredFile>,
252}
253
254impl OrganizedCode {
255 pub fn to_formatted_text(&self, include_content: bool, _root_path: &Path) -> Result<String> {
256 let mut output = String::new();
257
258 output.push_str(&format!(
259 r#"================================================================================
260 OMNIVORE CODE ANALYSIS REPORT
261================================================================================
262
263PROJECT INFORMATION
264-------------------
265Type: {}
266Description: {}
267Language: {}
268Frameworks: {}
269Build Tools: {}
270Total Files: {}
271
272================================================================================
273"#,
274 self.metadata.project_type,
275 self.metadata.description,
276 self.metadata.main_language,
277 if self.metadata.frameworks.is_empty() {
278 "None detected".to_string()
279 } else {
280 self.metadata.frameworks.join(", ")
281 },
282 if self.metadata.build_tools.is_empty() {
283 "None detected".to_string()
284 } else {
285 self.metadata.build_tools.join(", ")
286 },
287 self.metadata.total_files
288 ));
289
290 output.push_str("\nPROJECT STRUCTURE\n");
291 output.push_str("-----------------\n\n");
292
293 for section in &self.sections {
294 output.push_str(&format!("š {} ({})\n", section.name, section.files.len()));
295 output.push_str(&format!(" {}\n\n", section.description));
296
297 for file in §ion.files {
298 output.push_str(&format!(" ā¢ {}\n", file.relative_path.display()));
299 }
300 output.push_str("\n");
301 }
302
303 if include_content {
304 output.push_str("\n");
305 output.push_str("================================================================================\n");
306 output.push_str(" SOURCE CODE\n");
307 output.push_str("================================================================================\n\n");
308
309 for section in &self.sections {
310 if section.files.is_empty() {
311 continue;
312 }
313
314 output.push_str(&format!(
315 "\nāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāā\n"
316 ));
317 output.push_str(&format!(
318 "ā {} - {} file(s)\n",
319 section.name.to_uppercase(),
320 section.files.len()
321 ));
322 output.push_str(&format!(
323 "āāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāā\n\n"
324 ));
325
326 for file in §ion.files {
327 output.push_str(&format!(
328 "āāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāā\n"
329 ));
330 output.push_str(&format!("ā File: {}\n", file.relative_path.display()));
331 output.push_str(&format!(
332 "āāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāāā\n\n"
333 ));
334
335 if let Ok(content) = std::fs::read_to_string(&file.path) {
336 let lines: Vec<&str> = content.lines().collect();
337 for (i, line) in lines.iter().enumerate() {
338 output.push_str(&format!("{:4} ā {}\n", i + 1, line));
339 }
340 } else {
341 output.push_str("[Unable to read file content]\n");
342 }
343 output.push_str("\n");
344 }
345 }
346 }
347
348 output.push_str("\n================================================================================\n");
349 output.push_str(" Generated by Omnivore Code Extractor\n");
350 output.push_str("================================================================================\n");
351
352 Ok(output)
353 }
354
355 pub fn to_json(&self) -> Result<String> {
356 let mut json_output = serde_json::json!({
357 "metadata": {
358 "project_type": self.metadata.project_type,
359 "description": self.metadata.description,
360 "main_language": self.metadata.main_language,
361 "frameworks": self.metadata.frameworks,
362 "build_tools": self.metadata.build_tools,
363 "total_files": self.metadata.total_files,
364 },
365 "sections": []
366 });
367
368 let sections = json_output["sections"].as_array_mut().unwrap();
369
370 for section in &self.sections {
371 let mut section_json = serde_json::json!({
372 "name": section.name,
373 "description": section.description,
374 "files": []
375 });
376
377 let files = section_json["files"].as_array_mut().unwrap();
378 for file in §ion.files {
379 let content = std::fs::read_to_string(&file.path).unwrap_or_default();
380 files.push(serde_json::json!({
381 "path": file.relative_path.display().to_string(),
382 "content": content
383 }));
384 }
385
386 sections.push(section_json);
387 }
388
389 Ok(serde_json::to_string_pretty(&json_output)?)
390 }
391}