photoncloud-monorepo/creditservice/crates/creditservice-api/src/billing.rs

//! Billing module for CreditService
//!
//! Provides periodic billing functionality that charges projects based on usage metrics.

use async_trait::async_trait;
use chrono::{DateTime, Utc};
use creditservice_types::{ResourceType, Result};
use std::collections::HashMap;

/// Usage metrics for a project over a billing period
#[derive(Debug, Clone, Default)]
pub struct UsageMetrics {
    /// Project ID
    pub project_id: String,
    /// Resource usage by type (resource_type -> quantity)
    pub resource_usage: HashMap<ResourceType, ResourceUsage>,
    /// Billing period start
    pub period_start: DateTime<Utc>,
    /// Billing period end
    pub period_end: DateTime<Utc>,
}

/// Usage for a specific resource type
#[derive(Debug, Clone)]
pub struct ResourceUsage {
    /// Resource type
    pub resource_type: ResourceType,
    /// Total quantity used (e.g., VM-hours, GB-hours)
    pub quantity: f64,
    /// Unit for the quantity
    pub unit: String,
}

impl ResourceUsage {
    /// Create a new ResourceUsage
    pub fn new(resource_type: ResourceType, quantity: f64, unit: impl Into<String>) -> Self {
        Self {
            resource_type,
            quantity,
            unit: unit.into(),
        }
    }
}

/// Pricing rules for billing calculation
#[derive(Debug, Clone)]
pub struct PricingRules {
    /// Price per unit by resource type (resource_type -> credits per unit)
    pub prices: HashMap<ResourceType, i64>,
}

impl Default for PricingRules {
    fn default() -> Self {
        let mut prices = HashMap::new();
        // Default pricing (credits per hour/GB)
        prices.insert(ResourceType::VmInstance, 100);      // 100 credits/hour
        prices.insert(ResourceType::VmCpu, 10);            // 10 credits/CPU-hour
        prices.insert(ResourceType::VmMemoryGb, 5);        // 5 credits/GB-hour
        prices.insert(ResourceType::StorageGb, 1);         // 1 credit/GB-hour
        prices.insert(ResourceType::NetworkPort, 2);       // 2 credits/port-hour
        prices.insert(ResourceType::LoadBalancer, 50);     // 50 credits/hour
        prices.insert(ResourceType::DnsZone, 10);          // 10 credits/zone-hour
        prices.insert(ResourceType::DnsRecord, 1);         // 1 credit/record-hour
        prices.insert(ResourceType::K8sCluster, 200);      // 200 credits/hour
        prices.insert(ResourceType::K8sNode, 100);         // 100 credits/node-hour
        Self { prices }
    }
}

impl PricingRules {
    /// Calculate total charge for usage metrics
    pub fn calculate_charge(&self, usage: &UsageMetrics) -> i64 {
        let mut total: i64 = 0;
        for (resource_type, resource_usage) in &usage.resource_usage {
            if let Some(&price) = self.prices.get(resource_type) {
                // Calculate charge: quantity * price (rounded to nearest credit)
                let charge = (resource_usage.quantity * price as f64).round() as i64;
                total += charge;
            }
        }
        total
    }
}

/// Trait for fetching usage metrics (implemented by NightLight integration in S5)
#[async_trait]
pub trait UsageMetricsProvider: Send + Sync {
    /// Get usage metrics for a project over a billing period
    async fn get_usage_metrics(
        &self,
        project_id: &str,
        period_start: DateTime<Utc>,
        period_end: DateTime<Utc>,
    ) -> Result<UsageMetrics>;

    /// Get list of all projects with usage in the period
    async fn list_projects_with_usage(
        &self,
        period_start: DateTime<Utc>,
        period_end: DateTime<Utc>,
    ) -> Result<Vec<String>>;
}

/// Mock usage metrics provider for testing and until S5 is complete
#[derive(Debug, Default)]
pub struct MockUsageMetricsProvider {
    /// Predefined usage data for testing
    pub mock_data: HashMap<String, UsageMetrics>,
}

impl MockUsageMetricsProvider {
    /// Create a new mock provider
    pub fn new() -> Self {
        Self::default()
    }

    /// Add mock usage data for a project
    pub fn add_usage(&mut self, project_id: String, usage: UsageMetrics) {
        self.mock_data.insert(project_id, usage);
    }
}

#[async_trait]
impl UsageMetricsProvider for MockUsageMetricsProvider {
    async fn get_usage_metrics(
        &self,
        project_id: &str,
        period_start: DateTime<Utc>,
        period_end: DateTime<Utc>,
    ) -> Result<UsageMetrics> {
        Ok(self.mock_data.get(project_id).cloned().unwrap_or_else(|| UsageMetrics {
            project_id: project_id.to_string(),
            resource_usage: HashMap::new(),
            period_start,
            period_end,
        }))
    }

    async fn list_projects_with_usage(
        &self,
        _period_start: DateTime<Utc>,
        _period_end: DateTime<Utc>,
    ) -> Result<Vec<String>> {
        Ok(self.mock_data.keys().cloned().collect())
    }
}

/// Billing result for a single project
#[derive(Debug, Clone)]
pub struct ProjectBillingResult {
    pub project_id: String,
    pub amount_charged: i64,
    pub success: bool,
    pub error: Option<String>,
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_pricing_calculation() {
        let pricing = PricingRules::default();

        let mut usage = UsageMetrics::default();
        usage.resource_usage.insert(
            ResourceType::VmInstance,
            ResourceUsage::new(ResourceType::VmInstance, 10.0, "hours"),
        );
        usage.resource_usage.insert(
            ResourceType::StorageGb,
            ResourceUsage::new(ResourceType::StorageGb, 100.0, "GB-hours"),
        );

        let charge = pricing.calculate_charge(&usage);
        // 10 hours * 100 credits + 100 GB-hours * 1 credit = 1100 credits
        assert_eq!(charge, 1100);
    }

    #[tokio::test]
    async fn test_mock_usage_provider() {
        let mut provider = MockUsageMetricsProvider::new();

        let mut usage = UsageMetrics {
            project_id: "proj-1".into(),
            resource_usage: HashMap::new(),
            period_start: Utc::now(),
            period_end: Utc::now(),
        };
        usage.resource_usage.insert(
            ResourceType::VmInstance,
            ResourceUsage::new(ResourceType::VmInstance, 5.0, "hours"),
        );
        provider.add_usage("proj-1".into(), usage);

        let metrics = provider
            .get_usage_metrics("proj-1", Utc::now(), Utc::now())
            .await
            .unwrap();

        assert_eq!(metrics.project_id, "proj-1");
        assert!(metrics.resource_usage.contains_key(&ResourceType::VmInstance));
    }
}