1use rust_decimal::Decimal;
35use serde::{Deserialize, Serialize};
36use std::collections::HashMap;
37use uuid::Uuid;
38
39use crate::error::{AiError, Result};
40
41#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
43pub enum AccessTier {
44 Free = 0,
46 Bronze = 1,
48 Silver = 2,
50 Gold = 3,
52 Platinum = 4,
54}
55
56impl AccessTier {
57 #[must_use]
59 pub fn from_balance(balance: Decimal) -> Self {
60 if balance >= Decimal::new(100_000, 0) {
61 AccessTier::Platinum
62 } else if balance >= Decimal::new(10_000, 0) {
63 AccessTier::Gold
64 } else if balance >= Decimal::new(1_000, 0) {
65 AccessTier::Silver
66 } else if balance >= Decimal::ONE {
67 AccessTier::Bronze
68 } else {
69 AccessTier::Free
70 }
71 }
72
73 #[must_use]
75 pub fn min_balance(&self) -> Decimal {
76 match self {
77 AccessTier::Free => Decimal::ZERO,
78 AccessTier::Bronze => Decimal::ONE,
79 AccessTier::Silver => Decimal::new(1_000, 0),
80 AccessTier::Gold => Decimal::new(10_000, 0),
81 AccessTier::Platinum => Decimal::new(100_000, 0),
82 }
83 }
84
85 #[must_use]
87 pub fn name(&self) -> &'static str {
88 match self {
89 AccessTier::Free => "Free",
90 AccessTier::Bronze => "Bronze",
91 AccessTier::Silver => "Silver",
92 AccessTier::Gold => "Gold",
93 AccessTier::Platinum => "Platinum",
94 }
95 }
96
97 #[must_use]
99 pub fn description(&self) -> &'static str {
100 match self {
101 AccessTier::Free => "Basic access to AI features with limited usage",
102 AccessTier::Bronze => "Standard access with moderate usage limits",
103 AccessTier::Silver => "Enhanced access with higher limits and priority support",
104 AccessTier::Gold => "Premium access with very high limits and dedicated support",
105 AccessTier::Platinum => {
106 "Ultimate access with unlimited usage and personalized AI agents"
107 }
108 }
109 }
110}
111
112#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
114pub enum AiFeature {
115 CodeEvaluation,
117 ContentEvaluation,
119 CommitmentVerification,
121 FraudDetection,
123 GitHubVerification,
125 SocialMediaVerification,
127 DocumentAnalysis,
129 OcrAnalysis,
131 TranscriptAnalysis,
133 ReputationPrediction,
135 SentimentAnalysis,
137 MarketPrediction,
139 CustomAgent,
141}
142
143impl AiFeature {
144 #[must_use]
146 pub fn name(&self) -> &'static str {
147 match self {
148 AiFeature::CodeEvaluation => "Code Evaluation",
149 AiFeature::ContentEvaluation => "Content Evaluation",
150 AiFeature::CommitmentVerification => "Commitment Verification",
151 AiFeature::FraudDetection => "Fraud Detection",
152 AiFeature::GitHubVerification => "GitHub Verification",
153 AiFeature::SocialMediaVerification => "Social Media Verification",
154 AiFeature::DocumentAnalysis => "Document Analysis",
155 AiFeature::OcrAnalysis => "OCR Analysis",
156 AiFeature::TranscriptAnalysis => "Transcript Analysis",
157 AiFeature::ReputationPrediction => "Reputation Prediction",
158 AiFeature::SentimentAnalysis => "Sentiment Analysis",
159 AiFeature::MarketPrediction => "Market Prediction",
160 AiFeature::CustomAgent => "Custom AI Agent",
161 }
162 }
163}
164
165#[derive(Debug, Clone, Serialize, Deserialize)]
167pub struct FeatureQuota {
168 pub daily_limit: Option<u32>,
170 pub monthly_limit: Option<u32>,
172 pub daily_used: u32,
174 pub monthly_used: u32,
176 pub enabled: bool,
178}
179
180impl FeatureQuota {
181 #[must_use]
183 pub fn unlimited() -> Self {
184 Self {
185 daily_limit: None,
186 monthly_limit: None,
187 daily_used: 0,
188 monthly_used: 0,
189 enabled: true,
190 }
191 }
192
193 #[must_use]
195 pub fn limited(daily: u32, monthly: u32) -> Self {
196 Self {
197 daily_limit: Some(daily),
198 monthly_limit: Some(monthly),
199 daily_used: 0,
200 monthly_used: 0,
201 enabled: true,
202 }
203 }
204
205 #[must_use]
207 pub fn is_available(&self) -> bool {
208 if !self.enabled {
209 return false;
210 }
211 if let Some(daily) = self.daily_limit {
212 if self.daily_used >= daily {
213 return false;
214 }
215 }
216 if let Some(monthly) = self.monthly_limit {
217 if self.monthly_used >= monthly {
218 return false;
219 }
220 }
221 true
222 }
223
224 pub fn record_usage(&mut self) {
226 self.daily_used += 1;
227 self.monthly_used += 1;
228 }
229
230 pub fn reset_daily(&mut self) {
232 self.daily_used = 0;
233 }
234
235 pub fn reset_monthly(&mut self) {
237 self.monthly_used = 0;
238 self.daily_used = 0;
239 }
240}
241
242#[derive(Debug, Clone, Serialize, Deserialize)]
244pub struct TierConfig {
245 pub tier: AccessTier,
247 pub feature_quotas: HashMap<AiFeature, FeatureQuota>,
249 pub allow_custom_agents: bool,
251 pub max_custom_agents: u32,
253}
254
255impl TierConfig {
256 #[must_use]
258 pub fn free_tier() -> Self {
259 let mut quotas = HashMap::new();
260 quotas.insert(AiFeature::CodeEvaluation, FeatureQuota::limited(5, 50));
261 quotas.insert(AiFeature::ContentEvaluation, FeatureQuota::limited(5, 50));
262 quotas.insert(
263 AiFeature::CommitmentVerification,
264 FeatureQuota::limited(3, 20),
265 );
266 quotas.insert(AiFeature::FraudDetection, FeatureQuota::limited(3, 20));
267
268 Self {
269 tier: AccessTier::Free,
270 feature_quotas: quotas,
271 allow_custom_agents: false,
272 max_custom_agents: 0,
273 }
274 }
275
276 #[must_use]
278 pub fn bronze_tier() -> Self {
279 let mut quotas = HashMap::new();
280 quotas.insert(AiFeature::CodeEvaluation, FeatureQuota::limited(50, 500));
281 quotas.insert(AiFeature::ContentEvaluation, FeatureQuota::limited(50, 500));
282 quotas.insert(
283 AiFeature::CommitmentVerification,
284 FeatureQuota::limited(30, 300),
285 );
286 quotas.insert(AiFeature::FraudDetection, FeatureQuota::limited(30, 300));
287 quotas.insert(
288 AiFeature::GitHubVerification,
289 FeatureQuota::limited(20, 200),
290 );
291 quotas.insert(
292 AiFeature::SocialMediaVerification,
293 FeatureQuota::limited(20, 200),
294 );
295
296 Self {
297 tier: AccessTier::Bronze,
298 feature_quotas: quotas,
299 allow_custom_agents: false,
300 max_custom_agents: 0,
301 }
302 }
303
304 #[must_use]
306 pub fn silver_tier() -> Self {
307 let mut quotas = HashMap::new();
308 quotas.insert(AiFeature::CodeEvaluation, FeatureQuota::limited(200, 2000));
309 quotas.insert(
310 AiFeature::ContentEvaluation,
311 FeatureQuota::limited(200, 2000),
312 );
313 quotas.insert(
314 AiFeature::CommitmentVerification,
315 FeatureQuota::limited(100, 1000),
316 );
317 quotas.insert(AiFeature::FraudDetection, FeatureQuota::limited(100, 1000));
318 quotas.insert(
319 AiFeature::GitHubVerification,
320 FeatureQuota::limited(100, 1000),
321 );
322 quotas.insert(
323 AiFeature::SocialMediaVerification,
324 FeatureQuota::limited(100, 1000),
325 );
326 quotas.insert(AiFeature::DocumentAnalysis, FeatureQuota::limited(50, 500));
327 quotas.insert(AiFeature::OcrAnalysis, FeatureQuota::limited(50, 500));
328 quotas.insert(
329 AiFeature::TranscriptAnalysis,
330 FeatureQuota::limited(50, 500),
331 );
332
333 Self {
334 tier: AccessTier::Silver,
335 feature_quotas: quotas,
336 allow_custom_agents: false,
337 max_custom_agents: 0,
338 }
339 }
340
341 #[must_use]
343 pub fn gold_tier() -> Self {
344 let mut quotas = HashMap::new();
345 quotas.insert(
346 AiFeature::CodeEvaluation,
347 FeatureQuota::limited(1000, 10000),
348 );
349 quotas.insert(
350 AiFeature::ContentEvaluation,
351 FeatureQuota::limited(1000, 10000),
352 );
353 quotas.insert(
354 AiFeature::CommitmentVerification,
355 FeatureQuota::limited(500, 5000),
356 );
357 quotas.insert(AiFeature::FraudDetection, FeatureQuota::limited(500, 5000));
358 quotas.insert(
359 AiFeature::GitHubVerification,
360 FeatureQuota::limited(500, 5000),
361 );
362 quotas.insert(
363 AiFeature::SocialMediaVerification,
364 FeatureQuota::limited(500, 5000),
365 );
366 quotas.insert(
367 AiFeature::DocumentAnalysis,
368 FeatureQuota::limited(200, 2000),
369 );
370 quotas.insert(AiFeature::OcrAnalysis, FeatureQuota::limited(200, 2000));
371 quotas.insert(
372 AiFeature::TranscriptAnalysis,
373 FeatureQuota::limited(200, 2000),
374 );
375 quotas.insert(
376 AiFeature::ReputationPrediction,
377 FeatureQuota::limited(100, 1000),
378 );
379 quotas.insert(
380 AiFeature::SentimentAnalysis,
381 FeatureQuota::limited(100, 1000),
382 );
383
384 Self {
385 tier: AccessTier::Gold,
386 feature_quotas: quotas,
387 allow_custom_agents: true,
388 max_custom_agents: 1,
389 }
390 }
391
392 #[must_use]
394 pub fn platinum_tier() -> Self {
395 let mut quotas = HashMap::new();
396 quotas.insert(AiFeature::CodeEvaluation, FeatureQuota::unlimited());
398 quotas.insert(AiFeature::ContentEvaluation, FeatureQuota::unlimited());
399 quotas.insert(AiFeature::CommitmentVerification, FeatureQuota::unlimited());
400 quotas.insert(AiFeature::FraudDetection, FeatureQuota::unlimited());
401 quotas.insert(AiFeature::GitHubVerification, FeatureQuota::unlimited());
402 quotas.insert(
403 AiFeature::SocialMediaVerification,
404 FeatureQuota::unlimited(),
405 );
406 quotas.insert(AiFeature::DocumentAnalysis, FeatureQuota::unlimited());
407 quotas.insert(AiFeature::OcrAnalysis, FeatureQuota::unlimited());
408 quotas.insert(AiFeature::TranscriptAnalysis, FeatureQuota::unlimited());
409 quotas.insert(AiFeature::ReputationPrediction, FeatureQuota::unlimited());
410 quotas.insert(AiFeature::SentimentAnalysis, FeatureQuota::unlimited());
411 quotas.insert(AiFeature::MarketPrediction, FeatureQuota::unlimited());
412 quotas.insert(AiFeature::CustomAgent, FeatureQuota::unlimited());
413
414 Self {
415 tier: AccessTier::Platinum,
416 feature_quotas: quotas,
417 allow_custom_agents: true,
418 max_custom_agents: 5,
419 }
420 }
421
422 #[must_use]
424 pub fn get_quota(&self, feature: AiFeature) -> Option<&FeatureQuota> {
425 self.feature_quotas.get(&feature)
426 }
427
428 pub fn get_quota_mut(&mut self, feature: AiFeature) -> Option<&mut FeatureQuota> {
430 self.feature_quotas.get_mut(&feature)
431 }
432}
433
434#[derive(Debug, Clone, Serialize, Deserialize)]
436pub struct CustomAgentConfig {
437 pub agent_id: Uuid,
439 pub token_id: Uuid,
441 pub name: String,
443 pub description: Option<String>,
445 pub system_prompt: String,
447 pub model: String,
449 pub temperature: f32,
451 pub is_active: bool,
453 pub min_tier: AccessTier,
455}
456
457impl CustomAgentConfig {
458 #[must_use]
460 pub fn new(token_id: Uuid, name: String, system_prompt: String, model: String) -> Self {
461 Self {
462 agent_id: Uuid::new_v4(),
463 token_id,
464 name,
465 description: None,
466 system_prompt,
467 model,
468 temperature: 0.7,
469 is_active: true,
470 min_tier: AccessTier::Gold,
471 }
472 }
473}
474
475#[derive(Debug, Clone, Serialize, Deserialize)]
477pub struct TokenHolder {
478 pub user_id: Uuid,
480 pub token_id: Uuid,
482 pub balance: Decimal,
484 pub tier: AccessTier,
486}
487
488impl TokenHolder {
489 #[must_use]
491 pub fn new(user_id: Uuid, token_id: Uuid, balance: Decimal) -> Self {
492 let tier = AccessTier::from_balance(balance);
493 Self {
494 user_id,
495 token_id,
496 balance,
497 tier,
498 }
499 }
500}
501
502pub struct AccessControlManager {
504 tier_configs: HashMap<AccessTier, TierConfig>,
506 custom_agents: HashMap<Uuid, Vec<CustomAgentConfig>>,
508}
509
510impl Default for AccessControlManager {
511 fn default() -> Self {
512 Self::new()
513 }
514}
515
516impl AccessControlManager {
517 #[must_use]
519 pub fn new() -> Self {
520 let mut tier_configs = HashMap::new();
521 tier_configs.insert(AccessTier::Free, TierConfig::free_tier());
522 tier_configs.insert(AccessTier::Bronze, TierConfig::bronze_tier());
523 tier_configs.insert(AccessTier::Silver, TierConfig::silver_tier());
524 tier_configs.insert(AccessTier::Gold, TierConfig::gold_tier());
525 tier_configs.insert(AccessTier::Platinum, TierConfig::platinum_tier());
526
527 Self {
528 tier_configs,
529 custom_agents: HashMap::new(),
530 }
531 }
532
533 pub fn can_access_feature(&self, holder: &TokenHolder, feature: AiFeature) -> Result<bool> {
539 let config = self
540 .tier_configs
541 .get(&holder.tier)
542 .ok_or_else(|| AiError::Configuration("Invalid access tier".to_string()))?;
543
544 if let Some(quota) = config.get_quota(feature) {
545 Ok(quota.is_available())
546 } else {
547 Ok(false)
549 }
550 }
551
552 pub fn record_usage(&mut self, holder: &TokenHolder, feature: AiFeature) -> Result<()> {
561 let config = self
562 .tier_configs
563 .get_mut(&holder.tier)
564 .ok_or_else(|| AiError::Configuration("Invalid access tier".to_string()))?;
565
566 if let Some(quota) = config.get_quota_mut(feature) {
567 if !quota.is_available() {
568 return Err(AiError::QuotaExceeded(format!(
569 "Daily or monthly limit exceeded for {}",
570 feature.name()
571 )));
572 }
573 quota.record_usage();
574 Ok(())
575 } else {
576 Err(AiError::FeatureNotAvailable(format!(
577 "Feature {} not available for {} tier",
578 feature.name(),
579 holder.tier.name()
580 )))
581 }
582 }
583
584 #[must_use]
586 pub fn get_available_features(&self, tier: AccessTier) -> Vec<AiFeature> {
587 self.tier_configs
588 .get(&tier)
589 .map(|config| {
590 config
591 .feature_quotas
592 .iter()
593 .filter(|(_, quota)| quota.enabled)
594 .map(|(feature, _)| *feature)
595 .collect()
596 })
597 .unwrap_or_default()
598 }
599
600 pub fn register_custom_agent(
609 &mut self,
610 holder: &TokenHolder,
611 agent: CustomAgentConfig,
612 ) -> Result<()> {
613 let config = self
615 .tier_configs
616 .get(&holder.tier)
617 .ok_or_else(|| AiError::Configuration("Invalid access tier".to_string()))?;
618
619 if !config.allow_custom_agents {
620 return Err(AiError::FeatureNotAvailable(
621 "Custom agents not available for this tier".to_string(),
622 ));
623 }
624
625 let agents = self.custom_agents.entry(agent.token_id).or_default();
627 if agents.len() >= config.max_custom_agents as usize {
628 return Err(AiError::LimitExceeded(format!(
629 "Maximum {} custom agents allowed for {} tier",
630 config.max_custom_agents,
631 holder.tier.name()
632 )));
633 }
634
635 agents.push(agent);
636 Ok(())
637 }
638
639 #[must_use]
641 pub fn get_custom_agents(&self, token_id: Uuid) -> Vec<&CustomAgentConfig> {
642 self.custom_agents
643 .get(&token_id)
644 .map(|agents| agents.iter().collect())
645 .unwrap_or_default()
646 }
647
648 #[must_use]
650 pub fn get_tier_config(&self, tier: AccessTier) -> Option<&TierConfig> {
651 self.tier_configs.get(&tier)
652 }
653
654 pub fn update_tier_config(&mut self, config: TierConfig) {
656 self.tier_configs.insert(config.tier, config);
657 }
658}
659
660#[cfg(test)]
661mod tests {
662 use super::*;
663
664 #[test]
665 fn test_access_tier_from_balance() {
666 assert_eq!(AccessTier::from_balance(Decimal::ZERO), AccessTier::Free);
667 assert_eq!(
668 AccessTier::from_balance(Decimal::new(500, 0)),
669 AccessTier::Bronze
670 );
671 assert_eq!(
672 AccessTier::from_balance(Decimal::new(5000, 0)),
673 AccessTier::Silver
674 );
675 assert_eq!(
676 AccessTier::from_balance(Decimal::new(50000, 0)),
677 AccessTier::Gold
678 );
679 assert_eq!(
680 AccessTier::from_balance(Decimal::new(150000, 0)),
681 AccessTier::Platinum
682 );
683 }
684
685 #[test]
686 fn test_feature_quota_availability() {
687 let mut quota = FeatureQuota::limited(5, 50);
688 assert!(quota.is_available());
689
690 for _ in 0..5 {
692 quota.record_usage();
693 }
694 assert!(!quota.is_available());
695
696 quota.reset_daily();
698 assert!(quota.is_available());
699 }
700
701 #[test]
702 fn test_unlimited_quota() {
703 let mut quota = FeatureQuota::unlimited();
704 assert!(quota.is_available());
705
706 for _ in 0..1000 {
708 quota.record_usage();
709 }
710 assert!(quota.is_available());
711 }
712
713 #[test]
714 fn test_access_control_manager() {
715 let manager = AccessControlManager::new();
716
717 let holder = TokenHolder::new(
718 Uuid::new_v4(),
719 Uuid::new_v4(),
720 Decimal::new(5000, 0), );
722
723 assert!(
725 manager
726 .can_access_feature(&holder, AiFeature::DocumentAnalysis)
727 .unwrap()
728 );
729
730 assert!(
732 !manager
733 .can_access_feature(&holder, AiFeature::MarketPrediction)
734 .unwrap()
735 );
736 }
737
738 #[test]
739 fn test_custom_agent_registration() {
740 let mut manager = AccessControlManager::new();
741
742 let token_id = Uuid::new_v4();
743 let holder = TokenHolder::new(
744 Uuid::new_v4(),
745 token_id,
746 Decimal::new(50000, 0), );
748
749 let agent = CustomAgentConfig::new(
750 token_id,
751 "My Agent".to_string(),
752 "You are a helpful assistant".to_string(),
753 "gpt-4-turbo".to_string(),
754 );
755
756 assert!(manager.register_custom_agent(&holder, agent).is_ok());
758
759 let agent2 = CustomAgentConfig::new(
761 token_id,
762 "Agent 2".to_string(),
763 "Another prompt".to_string(),
764 "gpt-4-turbo".to_string(),
765 );
766 assert!(manager.register_custom_agent(&holder, agent2).is_err());
767
768 let agents = manager.get_custom_agents(token_id);
770 assert_eq!(agents.len(), 1);
771 }
772
773 #[test]
774 fn test_tier_ordering() {
775 assert!(AccessTier::Bronze > AccessTier::Free);
776 assert!(AccessTier::Silver > AccessTier::Bronze);
777 assert!(AccessTier::Gold > AccessTier::Silver);
778 assert!(AccessTier::Platinum > AccessTier::Gold);
779 }
780
781 #[test]
782 fn test_usage_tracking() {
783 let mut manager = AccessControlManager::new();
784
785 let holder = TokenHolder::new(
786 Uuid::new_v4(),
787 Uuid::new_v4(),
788 Decimal::new(100, 0), );
790
791 for _ in 0..5 {
793 assert!(
794 manager
795 .record_usage(&holder, AiFeature::CodeEvaluation)
796 .is_ok()
797 );
798 }
799
800 }
804}