1use crate::error::{Result, TokenError};
7use crate::treasury::FeeDistributionConfig;
8use dashmap::DashMap;
9use serde::{Deserialize, Serialize};
10use std::collections::HashMap;
11use tenzro_types::asset::AssetId;
12use tenzro_types::primitives::Timestamp;
13use tracing::{debug, info};
14
15#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
17pub enum FeeSource {
18 Transaction,
20 Settlement,
22 Bridge,
24 ModelInference,
26 Storage,
28 Other,
30}
31
32#[derive(Debug, Clone, Serialize, Deserialize)]
34pub struct FeeRecord {
35 pub fee_id: String,
37 pub asset_id: AssetId,
39 pub amount: u128,
41 pub source: FeeSource,
43 pub treasury_share: u128,
45 pub burn_share: u128,
47 pub staker_share: u128,
49 pub timestamp: Timestamp,
51}
52
53impl FeeRecord {
54 pub fn new(
56 asset_id: AssetId,
57 amount: u128,
58 source: FeeSource,
59 treasury_share: u128,
60 burn_share: u128,
61 staker_share: u128,
62 ) -> Self {
63 Self {
64 fee_id: uuid::Uuid::new_v4().to_string(),
65 asset_id,
66 amount,
67 source,
68 treasury_share,
69 burn_share,
70 staker_share,
71 timestamp: Timestamp::now(),
72 }
73 }
74}
75
76#[derive(Debug, Clone, Serialize, Deserialize)]
78#[derive(Default)]
79pub struct FeeStats {
80 pub total_collected: HashMap<AssetId, u128>,
82 pub by_source: HashMap<FeeSource, u128>,
84 pub total_to_treasury: u128,
86 pub total_burned: u128,
88 pub total_to_stakers: u128,
90 pub fee_count: u64,
92}
93
94
95#[derive(Debug, Clone, Serialize, Deserialize)]
97pub struct DistributionHistory {
98 pub period_start: Timestamp,
100 pub period_end: Timestamp,
102 pub total_fees: u128,
104 pub fees_by_asset: HashMap<AssetId, u128>,
106 pub treasury_share: u128,
108 pub burn_share: u128,
110 pub staker_share: u128,
112}
113
114impl DistributionHistory {
115 pub fn new(period_start: Timestamp, period_end: Timestamp) -> Self {
117 Self {
118 period_start,
119 period_end,
120 total_fees: 0,
121 fees_by_asset: HashMap::new(),
122 treasury_share: 0,
123 burn_share: 0,
124 staker_share: 0,
125 }
126 }
127
128 pub fn add_fee(&mut self, record: &FeeRecord) {
130 self.total_fees = self.total_fees.saturating_add(record.amount);
131
132 let asset_total = self.fees_by_asset.get(&record.asset_id).copied().unwrap_or(0);
133 self.fees_by_asset.insert(
134 record.asset_id.clone(),
135 asset_total.saturating_add(record.amount),
136 );
137
138 self.treasury_share = self.treasury_share.saturating_add(record.treasury_share);
139 self.burn_share = self.burn_share.saturating_add(record.burn_share);
140 self.staker_share = self.staker_share.saturating_add(record.staker_share);
141 }
142}
143
144#[derive(Debug)]
148pub struct FeeProcessor {
149 config: parking_lot::RwLock<FeeDistributionConfig>,
151 fee_records: DashMap<String, FeeRecord>,
153 stats: parking_lot::RwLock<FeeStats>,
155 history: DashMap<String, DistributionHistory>,
157 current_period: parking_lot::RwLock<Option<String>>,
159}
160
161impl FeeProcessor {
162 pub fn new() -> Self {
164 Self {
165 config: parking_lot::RwLock::new(FeeDistributionConfig::default()),
166 fee_records: DashMap::new(),
167 stats: parking_lot::RwLock::new(FeeStats::default()),
168 history: DashMap::new(),
169 current_period: parking_lot::RwLock::new(None),
170 }
171 }
172
173 pub fn process_fee(&self, asset_id: AssetId, amount: u128, source: FeeSource) -> Result<FeeRecord> {
183 if amount == 0 {
184 return Err(TokenError::InvalidAmount("Fee amount must be greater than zero".to_string()));
185 }
186
187 let config = self.config.read();
189 let distribution = config.calculate_distribution(amount);
190 drop(config);
191
192 let record = FeeRecord::new(
194 asset_id.clone(),
195 amount,
196 source,
197 distribution.treasury_amount,
198 distribution.burn_amount,
199 distribution.staker_amount,
200 );
201
202 let fee_id = record.fee_id.clone();
203
204 let mut stats = self.stats.write();
206
207 let asset_total = stats.total_collected.get(&asset_id).copied().unwrap_or(0);
209 stats.total_collected.insert(
210 asset_id.clone(),
211 asset_total.checked_add(amount)
212 .ok_or_else(|| TokenError::ArithmeticOverflow {
213 operation: "fee total collected".to_string(),
214 })?,
215 );
216
217 let source_total = stats.by_source.get(&source).copied().unwrap_or(0);
219 stats.by_source.insert(
220 source,
221 source_total.checked_add(amount)
222 .ok_or_else(|| TokenError::ArithmeticOverflow {
223 operation: "fee by source".to_string(),
224 })?,
225 );
226
227 stats.total_to_treasury = stats.total_to_treasury.checked_add(distribution.treasury_amount)
229 .ok_or_else(|| TokenError::ArithmeticOverflow {
230 operation: "fee total to treasury".to_string(),
231 })?;
232 stats.total_burned = stats.total_burned.checked_add(distribution.burn_amount)
233 .ok_or_else(|| TokenError::ArithmeticOverflow {
234 operation: "fee total burned".to_string(),
235 })?;
236 stats.total_to_stakers = stats.total_to_stakers.checked_add(distribution.staker_amount)
237 .ok_or_else(|| TokenError::ArithmeticOverflow {
238 operation: "fee total to stakers".to_string(),
239 })?;
240 stats.fee_count += 1;
241
242 drop(stats);
243
244 self.add_to_history(&record);
246
247 self.fee_records.insert(fee_id.clone(), record.clone());
249
250 debug!(
251 "Processed fee: {} of {} (treasury: {}, burn: {}, stakers: {})",
252 amount, asset_id.as_str(),
253 distribution.treasury_amount,
254 distribution.burn_amount,
255 distribution.staker_amount
256 );
257
258 Ok(record)
259 }
260
261 pub fn get_fee_stats(&self) -> FeeStats {
263 self.stats.read().clone()
264 }
265
266 pub fn get_distribution_history(&self, period: &str) -> Option<DistributionHistory> {
268 self.history.get(period).map(|h| h.clone())
269 }
270
271 pub fn get_all_history(&self) -> Vec<(String, DistributionHistory)> {
273 self.history
274 .iter()
275 .map(|entry| (entry.key().clone(), entry.value().clone()))
276 .collect()
277 }
278
279 pub fn start_new_period(&self, period_id: String, start_time: Timestamp) {
281 if let Some(current) = self.current_period.read().as_ref()
283 && let Some(mut history) = self.history.get_mut(current)
284 {
285 history.period_end = Timestamp::now();
286 }
287
288 let history = DistributionHistory::new(start_time, Timestamp::now());
290 self.history.insert(period_id.clone(), history);
291 *self.current_period.write() = Some(period_id.clone());
292
293 info!("Started new distribution period: {}", period_id);
294 }
295
296 pub fn update_config(&self, config: FeeDistributionConfig) -> Result<()> {
298 config.validate()?;
299 *self.config.write() = config;
300 info!("Updated fee distribution configuration");
301 Ok(())
302 }
303
304 pub fn get_config(&self) -> FeeDistributionConfig {
306 self.config.read().clone()
307 }
308
309 fn add_to_history(&self, record: &FeeRecord) {
311 let current_period = self.current_period.read().clone();
312
313 if let Some(period_id) = current_period {
314 if let Some(mut history) = self.history.get_mut(&period_id) {
315 history.add_fee(record);
316 }
317 } else {
318 drop(current_period);
320 let period_id = format!("period_{}", Timestamp::now().as_millis());
321 self.start_new_period(period_id.clone(), Timestamp::now());
322
323 if let Some(mut history) = self.history.get_mut(&period_id) {
324 history.add_fee(record);
325 }
326 }
327 }
328
329 pub fn get_fee_record(&self, fee_id: &str) -> Option<FeeRecord> {
331 self.fee_records.get(fee_id).map(|r| r.clone())
332 }
333
334 pub fn get_all_fee_records(&self) -> Vec<FeeRecord> {
336 self.fee_records
337 .iter()
338 .map(|entry| entry.value().clone())
339 .collect()
340 }
341}
342
343impl Default for FeeProcessor {
344 fn default() -> Self {
345 Self::new()
346 }
347}
348
349#[cfg(test)]
350mod tests {
351 use super::*;
352
353 #[test]
354 fn test_process_fee() {
355 let processor = FeeProcessor::new();
356 let asset_id = AssetId::tnzo();
357
358 let record = processor.process_fee(asset_id.clone(), 10000, FeeSource::Transaction).unwrap();
359
360 assert_eq!(record.amount, 10000);
361 assert_eq!(record.treasury_share, 4000); assert_eq!(record.burn_share, 3000); assert_eq!(record.staker_share, 3000); let stats = processor.get_fee_stats();
366 assert_eq!(stats.fee_count, 1);
367 assert_eq!(stats.total_to_treasury, 4000);
368 }
369
370 #[test]
371 fn test_fee_stats() {
372 let processor = FeeProcessor::new();
373 let asset_id = AssetId::tnzo();
374
375 processor.process_fee(asset_id.clone(), 10000, FeeSource::Transaction).unwrap();
376 processor.process_fee(asset_id.clone(), 5000, FeeSource::Settlement).unwrap();
377
378 let stats = processor.get_fee_stats();
379 assert_eq!(stats.fee_count, 2);
380 assert_eq!(*stats.total_collected.get(&asset_id).unwrap(), 15000);
381 }
382
383 #[test]
384 fn test_history_periods() {
385 let processor = FeeProcessor::new();
386 let asset_id = AssetId::tnzo();
387
388 processor.start_new_period("period1".to_string(), Timestamp::now());
389 processor.process_fee(asset_id.clone(), 10000, FeeSource::Transaction).unwrap();
390
391 let history = processor.get_distribution_history("period1").unwrap();
392 assert_eq!(history.total_fees, 10000);
393 }
394}