gitdataai/libs/agent/billing.rs

//! Billing service — handles user-level and project-level billing, deduction,
//! credit initialization, and error persistence.
//!
//! Architecture:
//!   - Each user gets $10 personal balance on signup.
//!   - Each project gets $20 balance only if it's the creator's first project,
//!     $0 otherwise.
//!   - AI usage is deducted from the project balance first; if insufficient,
//!     falls through to the user's personal balance.
//!   - Monthly quota only applies to pro users (is_pro = true).
//!   - If both project and user balance are insufficient, a billing_error
//!     record is persisted and an error is returned to the caller.

use db::database::AppDatabase;
use models::agents::model_pricing;
use models::ai::billing_error;
use models::projects::{project, project_billing, project_billing_history};
use models::users::user_billing;
use rust_decimal::Decimal;
use sea_orm::*;
use uuid::Uuid;

use crate::error::AgentError;

// ── Constants ──

fn default_user_balance() -> Decimal { Decimal::new(100_000, 4) } // $10.0000
fn first_project_credit() -> Decimal { Decimal::new(200_000, 4) } // $20.0000
const SUBSEQUENT_PROJECT_BALANCE: Decimal = Decimal::ZERO;

// ── Types ──

#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, utoipa::ToSchema)]
pub struct BillingRecord {
    pub cost: f64,
    pub currency: String,
    pub input_tokens: i64,
    pub output_tokens: i64,
    pub deducted_from: String, // "project" or "user"
}

#[derive(Debug)]
pub enum BillingResult {
    Success(BillingRecord),
    InsufficientBalance { message: String },
}

// ── Core deduction: AI usage ──

/// Record AI usage: deduct from project balance first, fall through to user balance.
///
/// Returns `InsufficientBalance` if neither account can cover the cost.
/// On insufficient balance, a `billing_error` record is persisted for frontend display.
pub async fn record_ai_usage(
    db: &AppDatabase,
    project_uid: Uuid,
    user_uid: Uuid,
    model_id: Uuid,
    input_tokens: i64,
    output_tokens: i64,
) -> Result<BillingResult, AgentError> {
    let total_cost = compute_cost(db, model_id, input_tokens, output_tokens).await?;
    let currency = get_currency(db, model_id).await?;

    // Verify project exists
    let _ = project::Entity::find_by_id(project_uid)
        .one(db)
        .await?
        .ok_or_else(|| AgentError::Internal("Project not found".into()))?;

    // Attempt project-level deduction first
    let project_result = deduct_from_project(db, project_uid, total_cost, &currency, model_id, input_tokens, output_tokens).await;

    match project_result {
        Ok(()) => {
            let cost_f64 = decimal_to_f64(total_cost);
            tracing::info!(
                project_id = %project_uid,
                model_id = %model_id,
                input_tokens, output_tokens,
                cost = %cost_f64,
                currency = %currency,
                deducted_from = "project",
                "ai_usage_recorded"
            );
            Ok(BillingResult::Success(BillingRecord {
                cost: cost_f64,
                currency,
                input_tokens,
                output_tokens,
                deducted_from: "project".to_string(),
            }))
        }
        Err(_) => {
            // Project balance insufficient — try user personal balance
            let user_result = deduct_from_user(db, user_uid, total_cost, &currency, project_uid, model_id, input_tokens, output_tokens).await;

            match user_result {
                Ok(()) => {
                    let cost_f64 = decimal_to_f64(total_cost);
                    tracing::info!(
                        user_id = %user_uid,
                        project_id = %project_uid,
                        model_id = %model_id,
                        input_tokens, output_tokens,
                        cost = %cost_f64,
                        currency = %currency,
                        deducted_from = "user",
                        "ai_usage_recorded"
                    );
                    Ok(BillingResult::Success(BillingRecord {
                        cost: cost_f64,
                        currency,
                        input_tokens,
                        output_tokens,
                        deducted_from: "user".to_string(),
                    }))
                }
                Err(insufficient_msg) => {
                    // Both project and user balance insufficient — persist error
                    persist_billing_error(
                        db,
                        "project",
                        project_uid,
                        "insufficient_balance",
                        &insufficient_msg,
                        Some(serde_json::json!({
                            "user_id": user_uid.to_string(),
                            "model_id": model_id.to_string(),
                            "input_tokens": input_tokens,
                            "output_tokens": output_tokens,
                            "cost": decimal_to_f64(total_cost),
                            "currency": currency,
                        })),
                    ).await?;

                    Ok(BillingResult::InsufficientBalance {
                        message: insufficient_msg,
                    })
                }
            }
        }
    }
}

/// Check whether a project + user has sufficient combined balance for a potential AI call.
/// Called before starting AI processing to avoid wasted compute.
pub async fn check_balance(
    db: &AppDatabase,
    project_uid: Uuid,
    user_uid: Uuid,
    model_id: Uuid,
    estimated_input_tokens: i64,
    estimated_output_tokens: i64,
) -> Result<bool, AgentError> {
    let estimated_cost = compute_cost(db, model_id, estimated_input_tokens, estimated_output_tokens).await?;
    let project_balance = get_project_balance(db, project_uid).await;
    let user_balance = get_user_balance(db, user_uid).await;

    Ok(project_balance + user_balance >= estimated_cost)
}

// ── Initialization ──

/// Initialize a user billing account with the default $10 balance.
/// Called on user signup / first login.
pub async fn initialize_user_billing(db: &AppDatabase, user_uid: Uuid) -> Result<(), AgentError> {
    let now = chrono::Utc::now();
    user_billing::ActiveModel {
        user: Set(user_uid),
        balance: Set(default_user_balance()),
        currency: Set("USD".to_string()),
        is_pro: Set(false),
        monthly_quota: Set(Decimal::ZERO),
        month_used: Set(Decimal::ZERO),
        cycle_start: Set(None),
        cycle_end: Set(None),
        updated_at: Set(now),
        created_at: Set(now),
    }
    .insert(db)
    .await
    .map_err(|e| AgentError::Internal(format!("failed to create user billing: {}", e)))?;

    tracing::info!(user_id = %user_uid, balance = "$10", "user_billing_initialized");
    Ok(())
}

/// Initialize a project billing account.
/// Grants $20 only if this is the creator's first project; $0 otherwise.
pub async fn initialize_project_billing(
    db: &AppDatabase,
    project_uid: Uuid,
    creator_uid: Uuid,
) -> Result<(), AgentError> {
    // Check how many projects this user has already created
    let existing_count = project::Entity::find()
        .filter(project::Column::CreatedBy.eq(creator_uid))
        .count(db)
        .await
        .map_err(|e| AgentError::Internal(format!("failed to count user projects: {}", e)))?;

    let is_first = existing_count == 0;
    let initial_balance = if is_first { first_project_credit() } else { SUBSEQUENT_PROJECT_BALANCE };
    let now = chrono::Utc::now();

    project_billing::ActiveModel {
        project: Set(project_uid),
        balance: Set(initial_balance),
        currency: Set("USD".to_string()),
        user: Set(Some(creator_uid)),
        initial_credit_granted: Set(is_first),
        is_pro: Set(false),
        monthly_quota: Set(Decimal::ZERO),
        month_used: Set(Decimal::ZERO),
        cycle_start: Set(None),
        cycle_end: Set(None),
        updated_at: Set(now),
        created_at: Set(now),
    }
    .insert(db)
    .await
    .map_err(|e| AgentError::Internal(format!("failed to create project billing: {}", e)))?;

    if is_first {
        // Record the credit in billing history
        project_billing_history::ActiveModel {
            uid: Set(Uuid::new_v4()),
            project: Set(project_uid),
            user: Set(Some(creator_uid)),
            amount: Set(first_project_credit()),
            currency: Set("USD".to_string()),
            reason: Set("first_project_credit".to_string()),
            extra: Set(Some(serde_json::json!({
                "is_first_project": true,
            }))),
            created_at: Set(now),
            ..Default::default()
        }
        .insert(db)
        .await
        .map_err(|e| AgentError::Internal(format!("failed to record credit history: {}", e)))?;
    }

    tracing::info!(
        project_id = %project_uid,
        creator_id = %creator_uid,
        is_first_project = is_first,
        balance = if is_first { "$20" } else { "$0" },
        "project_billing_initialized"
    );
    Ok(())
}

// ── Internal helpers ──

async fn compute_cost(
    db: &AppDatabase,
    model_id: Uuid,
    input_tokens: i64,
    output_tokens: i64,
) -> Result<Decimal, AgentError> {
    let pricing = model_pricing::Entity::find()
        .filter(model_pricing::Column::ModelVersionId.eq(model_id))
        .order_by_desc(model_pricing::Column::EffectiveFrom)
        .one(db)
        .await?
        .ok_or_else(|| AgentError::Internal(
            "No pricing record found for this model. Please configure AI model pricing first.".into(),
        ))?;

    let input_price: Decimal = pricing.input_price_per_1k_tokens.parse()
        .map_err(|e| AgentError::Internal(format!("Invalid input price: {}", e)))?;
    let output_price: Decimal = pricing.output_price_per_1k_tokens.parse()
        .map_err(|e| AgentError::Internal(format!("Invalid output price: {}", e)))?;

    let thousand = Decimal::from(1000);
    Ok((Decimal::from(input_tokens) / thousand) * input_price
        + (Decimal::from(output_tokens) / thousand) * output_price)
}

async fn get_currency(db: &AppDatabase, model_id: Uuid) -> Result<String, AgentError> {
    let pricing = model_pricing::Entity::find()
        .filter(model_pricing::Column::ModelVersionId.eq(model_id))
        .one(db)
        .await?
        .ok_or_else(|| AgentError::Internal("No pricing found".into()))?;
    Ok(pricing.currency.clone())
}

async fn get_project_balance(db: &AppDatabase, project_uid: Uuid) -> Decimal {
    project_billing::Entity::find_by_id(project_uid)
        .one(db)
        .await
        .ok()
        .flatten()
        .map(|b| b.balance)
        .unwrap_or(Decimal::ZERO)
}

async fn get_user_balance(db: &AppDatabase, user_uid: Uuid) -> Decimal {
    user_billing::Entity::find_by_id(user_uid)
        .one(db)
        .await
        .ok()
        .flatten()
        .map(|b| b.balance)
        .unwrap_or(Decimal::ZERO)
}

async fn deduct_from_project(
    db: &AppDatabase,
    project_uid: Uuid,
    cost: Decimal,
    currency: &str,
    model_id: Uuid,
    input_tokens: i64,
    output_tokens: i64,
) -> Result<(), String> {
    let txn = db.begin().await.map_err(|e| format!("db txn error: {}", e))?;

    let billing = project_billing::Entity::find_by_id(project_uid)
        .lock_exclusive()
        .one(&txn)
        .await
        .map_err(|e| format!("db error: {}", e))?
        .ok_or_else(|| "Project billing account not found".to_string())?;

    if billing.balance < cost {
        txn.rollback().await.ok();
        return Err(format!(
            "Project balance insufficient. Required: {:.4} {}, Available: {:.4} {}",
            cost, currency, billing.balance, currency
        ));
    }

    let now = chrono::Utc::now();

    project_billing_history::ActiveModel {
        uid: Set(Uuid::new_v4()),
        project: Set(project_uid),
        user: Set(None),
        amount: Set(-cost),
        currency: Set(currency.to_string()),
        reason: Set("ai_usage".to_string()),
        extra: Set(Some(serde_json::json!({
            "model_id": model_id.to_string(),
            "input_tokens": input_tokens,
            "output_tokens": output_tokens,
            "deducted_from": "project",
        }))),
        created_at: Set(now),
        ..Default::default()
    }
    .insert(&txn)
    .await
    .map_err(|e| format!("failed to insert history: {}", e))?;

    let mut updated: project_billing::ActiveModel = billing.into();
    updated.balance = Set(updated.balance.unwrap() - cost);
    updated.updated_at = Set(now);
    updated.update(&txn).await.map_err(|e| format!("failed to update balance: {}", e))?;

    txn.commit().await.map_err(|e| format!("commit error: {}", e))?;
    Ok(())
}

async fn deduct_from_user(
    db: &AppDatabase,
    user_uid: Uuid,
    cost: Decimal,
    currency: &str,
    project_uid: Uuid,
    model_id: Uuid,
    input_tokens: i64,
    output_tokens: i64,
) -> Result<(), String> {
    let txn = db.begin().await.map_err(|e| format!("db txn error: {}", e))?;

    let billing = user_billing::Entity::find_by_id(user_uid)
        .lock_exclusive()
        .one(&txn)
        .await
        .map_err(|e| format!("db error: {}", e))?
        .ok_or_else(|| "User billing account not found".to_string())?;

    if billing.balance < cost {
        txn.rollback().await.ok();
        return Err(format!(
            "Insufficient balance (project + user). Project: unavailable, User: {:.4} {}. Required: {:.4} {}",
            billing.balance, currency, cost, currency
        ));
    }

    let now = chrono::Utc::now();

    // Record in project billing history (but deducted from user)
    project_billing_history::ActiveModel {
        uid: Set(Uuid::new_v4()),
        project: Set(project_uid),
        user: Set(Some(user_uid)),
        amount: Set(-cost),
        currency: Set(currency.to_string()),
        reason: Set("ai_usage_user_fallback".to_string()),
        extra: Set(Some(serde_json::json!({
            "model_id": model_id.to_string(),
            "input_tokens": input_tokens,
            "output_tokens": output_tokens,
            "deducted_from": "user",
        }))),
        created_at: Set(now),
        ..Default::default()
    }
    .insert(&txn)
    .await
    .map_err(|e| format!("failed to insert history: {}", e))?;

    let mut updated: user_billing::ActiveModel = billing.into();
    updated.balance = Set(updated.balance.unwrap() - cost);
    updated.updated_at = Set(now);
    updated.update(&txn).await.map_err(|e| format!("failed to update user balance: {}", e))?;

    txn.commit().await.map_err(|e| format!("commit error: {}", e))?;
    Ok(())
}

pub async fn persist_billing_error(
    db: &AppDatabase,
    scope: &str,
    scope_id: Uuid,
    error_type: &str,
    message: &str,
    details: Option<serde_json::Value>,
) -> Result<(), AgentError> {
    billing_error::ActiveModel {
        id: Set(Uuid::new_v4()),
        scope: Set(scope.to_string()),
        scope_id: Set(scope_id),
        error_type: Set(error_type.to_string()),
        message: Set(message.to_string()),
        details: Set(details),
        resolved: Set(false),
        created_at: Set(chrono::Utc::now()),
    }
    .insert(db)
    .await
    .map_err(|e| AgentError::Internal(format!("failed to persist billing error: {}", e)))?;

    tracing::warn!(scope, %scope_id, error_type, "billing_error_persisted");
    Ok(())
}

fn decimal_to_f64(d: Decimal) -> f64 {
    d.round_dp(10).to_string().parse().unwrap_or(0.0)
}