1pub mod gain_score;
2pub mod model_pricing;
3pub mod task_classifier;
4
5use serde::{Deserialize, Serialize};
6
7use crate::core::a2a::cost_attribution::CostStore;
8use crate::core::gain::gain_score::GainScore;
9use crate::core::gain::model_pricing::{ModelPricing, ModelQuote};
10use crate::core::gain::task_classifier::{TaskCategory, TaskClassifier};
11use crate::core::heatmap::HeatMap;
12use crate::core::stats::StatsStore;
13
14#[derive(Clone)]
15pub struct GainEngine {
16 pub stats: StatsStore,
17 pub costs: CostStore,
18 pub heatmap: HeatMap,
19 pub pricing: ModelPricing,
20 pub events: Vec<crate::core::events::LeanCtxEvent>,
21 pub session: Option<crate::core::session::SessionState>,
22}
23
24#[derive(Debug, Clone, Serialize, Deserialize)]
25pub struct GainSummary {
26 pub model: ModelQuote,
27 pub total_commands: u64,
28 pub input_tokens: u64,
29 pub output_tokens: u64,
30 pub tokens_saved: u64,
31 pub gain_rate_pct: f64,
32 pub avoided_usd: f64,
33 pub energy_wh: f64,
35 pub co2_grams: f64,
37 pub tool_spend_usd: f64,
38 pub roi: Option<f64>,
39 pub score: GainScore,
40 #[serde(skip_serializing_if = "Option::is_none", default)]
41 pub daemon_hint: Option<String>,
42}
43
44#[derive(Debug, Clone, Serialize, Deserialize)]
45pub struct TaskGainRow {
46 pub category: TaskCategory,
47 pub commands: u64,
48 pub tokens_saved: u64,
49 pub tool_calls: u64,
50 pub tool_spend_usd: f64,
51}
52
53#[derive(Debug, Clone, Serialize, Deserialize)]
54pub struct FileGainRow {
55 pub path: String,
56 pub access_count: u32,
57 pub tokens_saved: u64,
58 pub compression_pct: f32,
59}
60
61impl GainEngine {
62 pub fn load() -> Self {
63 Self {
64 stats: crate::core::stats::load(),
65 costs: crate::core::a2a::cost_attribution::CostStore::load(),
66 heatmap: crate::core::heatmap::HeatMap::load(),
67 pricing: ModelPricing::load(),
68 events: crate::core::events::load_events_from_file(500),
69 session: crate::core::session::SessionState::load_latest(),
70 }
71 }
72
73 pub fn summary(&self, model: Option<&str>) -> GainSummary {
74 let quote = self.pricing.quote(model);
75 let tokens_saved = self
76 .stats
77 .total_input_tokens
78 .saturating_sub(self.stats.total_output_tokens);
79 let gain_rate_pct = if self.stats.total_input_tokens > 0 {
80 tokens_saved as f64 / self.stats.total_input_tokens as f64 * 100.0
81 } else {
82 0.0
83 };
84 let avoided_usd = quote.cost.estimate_usd(tokens_saved, 0, 0, 0);
85 let tool_spend_usd = self.costs.total_cost().max(0.0);
86 let roi = if tool_spend_usd > 0.0 {
87 Some(avoided_usd / tool_spend_usd)
88 } else {
89 None
90 };
91 let score = GainScore::compute(&self.stats, &self.costs, &self.pricing, model);
92 #[cfg(unix)]
93 let daemon_hint = if crate::daemon::is_daemon_running() {
94 None
95 } else {
96 Some(
97 "daemon not running — stats tracked locally (lean-ctx serve -d for full tracking)"
98 .to_string(),
99 )
100 };
101 #[cfg(not(unix))]
102 let daemon_hint: Option<String> = None;
103 GainSummary {
104 model: quote,
105 total_commands: self.stats.total_commands,
106 input_tokens: self.stats.total_input_tokens,
107 output_tokens: self.stats.total_output_tokens,
108 tokens_saved,
109 gain_rate_pct,
110 avoided_usd,
111 energy_wh: crate::core::energy::wh_for_tokens(tokens_saved),
112 co2_grams: crate::core::energy::co2_grams_for_tokens(tokens_saved),
113 tool_spend_usd,
114 roi,
115 score,
116 daemon_hint,
117 }
118 }
119
120 pub fn gain_score(&self, model: Option<&str>) -> GainScore {
121 GainScore::compute(&self.stats, &self.costs, &self.pricing, model)
122 }
123
124 pub fn task_breakdown(&self) -> Vec<TaskGainRow> {
125 use std::collections::HashMap;
126
127 let mut by_cat: HashMap<TaskCategory, TaskGainRow> = HashMap::new();
128
129 for (cmd_key, st) in &self.stats.commands {
130 let cat = TaskClassifier::classify_command_key(cmd_key);
131 let row = by_cat.entry(cat).or_insert(TaskGainRow {
132 category: cat,
133 commands: 0,
134 tokens_saved: 0,
135 tool_calls: 0,
136 tool_spend_usd: 0.0,
137 });
138 row.commands += st.count;
139 row.tokens_saved += st.input_tokens.saturating_sub(st.output_tokens);
140 }
141
142 for (tool, tc) in &self.costs.tools {
143 let cat = TaskClassifier::classify_tool(tool);
144 let row = by_cat.entry(cat).or_insert(TaskGainRow {
145 category: cat,
146 commands: 0,
147 tokens_saved: 0,
148 tool_calls: 0,
149 tool_spend_usd: 0.0,
150 });
151 row.tool_calls += tc.total_calls;
152 row.tool_spend_usd += tc.cost_usd;
153 }
154
155 let mut out: Vec<TaskGainRow> = by_cat.into_values().collect();
156 out.sort_by_key(|x| std::cmp::Reverse(x.tokens_saved));
157 out
158 }
159
160 pub fn heatmap_gains(&self, limit: usize) -> Vec<FileGainRow> {
161 let mut items: Vec<_> = self.heatmap.entries.values().collect();
162 items.sort_by_key(|x| std::cmp::Reverse(x.total_tokens_saved));
163 items.truncate(limit);
164 items
165 .into_iter()
166 .map(|e| FileGainRow {
167 path: e.path.clone(),
168 access_count: e.access_count,
169 tokens_saved: e.total_tokens_saved,
170 compression_pct: e.avg_compression_ratio * 100.0,
171 })
172 .collect()
173 }
174}