Skip to main content

st/formatters/
ai_json.rs

1use super::{ai::AiFormatter, Formatter, PathDisplayMode};
2use crate::scanner::{FileNode, TreeStats};
3use anyhow::Result;
4use serde_json::{json, Value};
5use std::io::Write;
6use std::path::Path;
7
8pub struct AiJsonFormatter {
9    ai_formatter: AiFormatter,
10}
11
12impl AiJsonFormatter {
13    pub fn new(no_emoji: bool, _path_mode: PathDisplayMode) -> Self {
14        Self {
15            ai_formatter: AiFormatter::new(no_emoji, _path_mode),
16        }
17    }
18}
19
20impl Formatter for AiJsonFormatter {
21    fn format(
22        &self,
23        writer: &mut dyn Write,
24        nodes: &[FileNode],
25        stats: &TreeStats,
26        root_path: &Path,
27    ) -> Result<()> {
28        // First get the AI format output as a string
29        let mut ai_output = Vec::new();
30        self.ai_formatter
31            .format(&mut ai_output, nodes, stats, root_path)?;
32        let ai_text = String::from_utf8_lossy(&ai_output);
33
34        // Parse the AI output to extract structured data
35        let lines = ai_text.lines();
36        let mut hex_tree_lines = Vec::new();
37        let mut context = None;
38        let mut hash = None;
39        let mut stats_section = false;
40        let mut file_count = 0u64;
41        let mut dir_count = 0u64;
42        let mut total_size = 0u64;
43        let mut file_types = Vec::new();
44        let mut large_files = Vec::new();
45        let mut date_range = None;
46
47        for line in lines {
48            if line == "TREE_HEX_V1:" {
49                continue;
50            } else if line.starts_with("CONTEXT: ") {
51                context = Some(line.strip_prefix("CONTEXT: ").unwrap_or("").to_string());
52            } else if line.starts_with("HASH: ") {
53                hash = Some(line.strip_prefix("HASH: ").unwrap_or("").to_string());
54            } else if line == "STATS:" || line.is_empty() {
55                stats_section = true;
56            } else if line.starts_with("F:") && stats_section {
57                // Parse stats line: F:45 D:12 S:23fc00 (2.3MB)
58                let parts: Vec<&str> = line.split_whitespace().collect();
59                if let Some(f) = parts.first().and_then(|s| s.strip_prefix("F:")) {
60                    file_count = u64::from_str_radix(f, 16).unwrap_or(0);
61                }
62                if let Some(d) = parts.get(1).and_then(|s| s.strip_prefix("D:")) {
63                    dir_count = u64::from_str_radix(d, 16).unwrap_or(0);
64                }
65                if let Some(s) = parts.get(2).and_then(|s| s.strip_prefix("S:")) {
66                    total_size = u64::from_str_radix(s, 16).unwrap_or(0);
67                }
68            } else if line.starts_with("TYPES: ") && stats_section {
69                let types_str = line.strip_prefix("TYPES: ").unwrap_or("");
70                for type_entry in types_str.split_whitespace() {
71                    if let Some((ext, count_hex)) = type_entry.split_once(':') {
72                        if let Ok(count) = u64::from_str_radix(count_hex, 16) {
73                            file_types.push(json!({
74                                "extension": ext,
75                                "count": count
76                            }));
77                        }
78                    }
79                }
80            } else if line.starts_with("LARGE: ") && stats_section {
81                let large_str = line.strip_prefix("LARGE: ").unwrap_or("");
82                for file_entry in large_str.split_whitespace() {
83                    if let Some((name, size_hex)) = file_entry.split_once(':') {
84                        if let Ok(size) = u64::from_str_radix(size_hex, 16) {
85                            large_files.push(json!({
86                                "name": name,
87                                "size": size
88                            }));
89                        }
90                    }
91                }
92            } else if line.starts_with("DATES: ") && stats_section {
93                date_range = Some(line.strip_prefix("DATES: ").unwrap_or("").to_string());
94            } else if line == "END_AI" {
95                break;
96            } else if !stats_section && !line.is_empty() {
97                // This is a hex tree line
98                hex_tree_lines.push(line.to_string());
99            }
100        }
101
102        // Build the JSON structure
103        let mut json_output = json!({
104            "version": "AI_JSON_V1",
105            "hash": hash.unwrap_or_else(|| "unknown".to_string()),
106            "hex_tree": hex_tree_lines,
107            "statistics": {
108                "files": file_count,
109                "directories": dir_count,
110                "total_size": total_size,
111                "total_size_mb": format!("{:.1}", total_size as f64 / (1024.0 * 1024.0))
112            }
113        });
114
115        // Add optional fields
116        if let Some(ctx) = context {
117            json_output["context"] = Value::String(ctx);
118        }
119
120        if !file_types.is_empty() {
121            json_output["statistics"]["file_types"] = Value::Array(file_types);
122        }
123
124        if !large_files.is_empty() {
125            json_output["statistics"]["largest_files"] = Value::Array(large_files);
126        }
127
128        if let Some(dates) = date_range {
129            json_output["statistics"]["date_range"] = Value::String(dates);
130        }
131
132        // Write the JSON output
133        writeln!(writer, "{}", serde_json::to_string_pretty(&json_output)?)?;
134
135        Ok(())
136    }
137}