1use crate::core::graph_provider::{self, GraphProvider};
2use std::collections::HashMap;
3
4struct HeatEntry {
5 path: String,
6 token_count: usize,
7 connections: usize,
8 heat_score: f64,
9}
10
11pub fn cmd_heatmap(args: &[String]) {
12 let project_root = std::env::current_dir()
13 .ok()
14 .and_then(|d| d.to_str().map(String::from))
15 .unwrap_or_else(|| ".".to_string());
16
17 let top_n: usize = args
18 .iter()
19 .find_map(|a| a.strip_prefix("--top="))
20 .and_then(|v| v.parse().ok())
21 .unwrap_or(20);
22
23 let dir_filter: Option<&str> = args
24 .iter()
25 .find_map(|a| a.strip_prefix("--dir="))
26 .map(|s| s.trim_end_matches('/'));
27
28 let sort_by = if args.iter().any(|a| a == "--by=connections") {
29 SortBy::Connections
30 } else if args.iter().any(|a| a == "--by=tokens") {
31 SortBy::Tokens
32 } else {
33 SortBy::Heat
34 };
35
36 let json_output = args.iter().any(|a| a == "--json");
37
38 let Some(open) = graph_provider::open_or_build(&project_root) else {
39 eprintln!("No graph available for project.");
40 return;
41 };
42
43 let entries = build_heat_entries(&open.provider, dir_filter);
44
45 if entries.is_empty() {
46 eprintln!("No files found in project graph.");
47 eprintln!(" Run: lean-ctx setup (to build the project graph)");
48 return;
49 }
50
51 let mut sorted = entries;
52 match sort_by {
53 SortBy::Heat => sorted.sort_by(|a, b| {
54 b.heat_score
55 .partial_cmp(&a.heat_score)
56 .unwrap_or(std::cmp::Ordering::Equal)
57 }),
58 SortBy::Tokens => sorted.sort_by_key(|x| std::cmp::Reverse(x.token_count)),
59 SortBy::Connections => sorted.sort_by_key(|x| std::cmp::Reverse(x.connections)),
60 }
61
62 let top = &sorted[..sorted.len().min(top_n)];
63
64 if json_output {
65 print_json(top);
66 } else {
67 print_heatmap(&project_root, top, &sorted);
68 }
69}
70
71enum SortBy {
72 Heat,
73 Tokens,
74 Connections,
75}
76
77fn build_heat_entries(gp: &GraphProvider, dir_filter: Option<&str>) -> Vec<HeatEntry> {
78 let all_edges = gp.edges();
79 let mut connection_counts: HashMap<String, usize> = HashMap::new();
80 for edge in &all_edges {
81 *connection_counts.entry(edge.from.clone()).or_default() += 1;
82 *connection_counts.entry(edge.to.clone()).or_default() += 1;
83 }
84
85 let paths = gp.file_paths();
86 let mut max_tokens = 1usize;
87 let mut file_entries: Vec<(String, usize)> = Vec::new();
88 for path in &paths {
89 if let Some(entry) = gp.get_file_entry(path) {
90 max_tokens = max_tokens.max(entry.token_count);
91 file_entries.push((path.clone(), entry.token_count));
92 }
93 }
94 let max_tokens = max_tokens as f64;
95 let max_connections = connection_counts.values().max().copied().unwrap_or(1) as f64;
96
97 file_entries
98 .into_iter()
99 .filter(|(path, _)| {
100 if let Some(dir) = dir_filter {
101 path.starts_with(dir) || path.starts_with(&format!("./{dir}"))
102 } else {
103 true
104 }
105 })
106 .map(|(path, token_count)| {
107 let connections = connection_counts.get(&path).copied().unwrap_or(0);
108 let token_norm = token_count as f64 / max_tokens;
109 let conn_norm = connections as f64 / max_connections;
110 let heat_score = token_norm * 0.4 + conn_norm * 0.6;
111
112 HeatEntry {
113 path,
114 token_count,
115 connections,
116 heat_score,
117 }
118 })
119 .collect()
120}
121
122fn heat_color(score: f64) -> &'static str {
123 if score > 0.8 {
124 "\x1b[91m" } else if score > 0.6 {
126 "\x1b[31m" } else if score > 0.4 {
128 "\x1b[33m" } else if score > 0.2 {
130 "\x1b[36m" } else {
132 "\x1b[34m" }
134}
135
136fn heat_bar(score: f64, width: usize) -> String {
137 let filled = (score * width as f64).round() as usize;
138 let blocks = "█".repeat(filled);
139 let empty = "░".repeat(width.saturating_sub(filled));
140 format!("{}{blocks}\x1b[38;5;239m{empty}\x1b[0m", heat_color(score))
141}
142
143fn print_heatmap(project_root: &str, entries: &[HeatEntry], all: &[HeatEntry]) {
144 let total_files = all.len();
145 let total_tokens: usize = all.iter().map(|e| e.token_count).sum();
146 let total_connections: usize = all.iter().map(|e| e.connections).sum();
147
148 let project_name = std::path::Path::new(project_root).file_name().map_or_else(
149 || project_root.to_string(),
150 |n| n.to_string_lossy().to_string(),
151 );
152
153 println!();
154 println!("\x1b[1;37m Context Heat Map\x1b[0m \x1b[38;5;239m{project_name}\x1b[0m");
155 println!(
156 "\x1b[38;5;239m {total_files} files · {total_tokens} tokens · {total_connections} connections\x1b[0m"
157 );
158 println!();
159
160 let max_path_len = entries.iter().map(|e| e.path.len()).max().unwrap_or(30);
161 let path_width = max_path_len.min(50);
162
163 println!(
164 " \x1b[38;5;239m{:<width$} {:>6} {:>5} HEAT\x1b[0m",
165 "FILE",
166 "TOKENS",
167 "CONNS",
168 width = path_width
169 );
170 println!(" \x1b[38;5;239m{}\x1b[0m", "─".repeat(path_width + 32));
171
172 for entry in entries {
173 let display_path = if entry.path.len() > path_width {
174 let skip = entry.path.len() - path_width + 3;
175 format!("...{}", &entry.path[skip..])
176 } else {
177 entry.path.clone()
178 };
179
180 let bar = heat_bar(entry.heat_score, 16);
181
182 println!(
183 " {color}{:<width$}\x1b[0m \x1b[38;5;245m{:>6}\x1b[0m \x1b[38;5;245m{:>5}\x1b[0m {bar} {color}{:.0}%\x1b[0m",
184 display_path,
185 entry.token_count,
186 entry.connections,
187 entry.heat_score * 100.0,
188 color = heat_color(entry.heat_score),
189 width = path_width,
190 );
191 }
192
193 println!();
194 println!(
195 " \x1b[38;5;239mLegend: \x1b[91m█\x1b[38;5;239m hot \x1b[33m█\x1b[38;5;239m warm \x1b[36m█\x1b[38;5;239m cool \x1b[34m█\x1b[38;5;239m cold\x1b[0m"
196 );
197 println!(
198 " \x1b[38;5;239mOptions: --top=N --dir=path --by=tokens|connections --json\x1b[0m"
199 );
200 println!();
201}
202
203fn print_json(entries: &[HeatEntry]) {
204 let items: Vec<serde_json::Value> = entries
205 .iter()
206 .map(|e| {
207 serde_json::json!({
208 "path": e.path,
209 "token_count": e.token_count,
210 "connections": e.connections,
211 "heat_score": (e.heat_score * 100.0).round() / 100.0,
212 })
213 })
214 .collect();
215
216 println!(
217 "{}",
218 serde_json::to_string_pretty(&items).unwrap_or_else(|_| "[]".to_string())
219 );
220}
221
222#[cfg(test)]
223mod tests {
224 use super::*;
225
226 #[test]
227 fn test_heat_color_ranges() {
228 assert_eq!(heat_color(0.9), "\x1b[91m");
229 assert_eq!(heat_color(0.7), "\x1b[31m");
230 assert_eq!(heat_color(0.5), "\x1b[33m");
231 assert_eq!(heat_color(0.3), "\x1b[36m");
232 assert_eq!(heat_color(0.1), "\x1b[34m");
233 }
234
235 #[test]
236 fn test_heat_bar_length() {
237 let bar = heat_bar(0.5, 10);
238 assert!(bar.contains("█████"));
239 }
240
241 #[test]
242 fn test_build_heat_entries_empty() {
243 let gp = GraphProvider::GraphIndex(crate::core::graph_index::ProjectIndex::new("."));
244 let entries = build_heat_entries(&gp, None);
245 assert!(entries.is_empty());
246 }
247}