gitdataai/libs/service/workspace/billing.rs

use crate::AppService;
use crate::error::AppError;
use chrono::{DateTime, Datelike, NaiveDate, Utc};
use models::Decimal;
use models::workspaces::{workspace_billing, workspace_billing_history, workspace_membership};
use sea_orm::sea_query::prelude::rust_decimal::prelude::ToPrimitive;
use sea_orm::*;
use serde::{Deserialize, Serialize};
use session::Session;
use utoipa::{IntoParams, ToSchema};
use uuid::Uuid;

/// Default monthly AI quota for workspace (shared across all its projects).
const DEFAULT_MONTHLY_QUOTA: f64 = 100.0;

#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
pub struct WorkspaceBillingCurrentResponse {
    pub workspace_id: Uuid,
    pub currency: String,
    pub monthly_quota: f64,
    pub balance: f64,
    pub total_spent: f64,
    pub month_used: f64,
    pub cycle_start_utc: DateTime<Utc>,
    pub cycle_end_utc: DateTime<Utc>,
    pub updated_at: DateTime<Utc>,
    pub created_at: DateTime<Utc>,
}

#[derive(Debug, Clone, Serialize, Deserialize, ToSchema, IntoParams)]
pub struct WorkspaceBillingHistoryQuery {
    pub page: Option<u64>,
    pub per_page: Option<u64>,
}

#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
pub struct WorkspaceBillingHistoryItem {
    pub uid: Uuid,
    pub workspace_id: Uuid,
    pub user_id: Option<Uuid>,
    pub amount: f64,
    pub currency: String,
    pub reason: String,
    pub extra: Option<serde_json::Value>,
    pub created_at: DateTime<Utc>,
}

#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
pub struct WorkspaceBillingHistoryResponse {
    pub page: u64,
    pub per_page: u64,
    pub total: u64,
    pub list: Vec<WorkspaceBillingHistoryItem>,
}

#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
pub struct WorkspaceBillingAddCreditParams {
    pub amount: f64,
    pub reason: Option<String>,
}

impl AppService {
    /// Get current workspace billing info.
    pub async fn workspace_billing_current(
        &self,
        ctx: &Session,
        workspace_slug: String,
    ) -> Result<WorkspaceBillingCurrentResponse, AppError> {
        let user_uid = ctx.user().ok_or(AppError::Unauthorized)?;
        let ws = self
            .utils_find_workspace_by_slug(workspace_slug.clone())
            .await?;
        let _ = workspace_membership::Entity::find()
            .filter(workspace_membership::Column::WorkspaceId.eq(ws.id))
            .filter(workspace_membership::Column::UserId.eq(user_uid))
            .filter(workspace_membership::Column::Status.eq("active"))
            .one(&self.db)
            .await?
            .ok_or(AppError::NotWorkspaceMember)?;

        let billing = self.ensure_workspace_billing(ws.id, Some(user_uid)).await?;
        let now_utc = Utc::now();
        let (month_start, next_month_start) = utc_month_bounds(now_utc)?;

        let month_used = workspace_billing_history::Entity::find()
            .filter(workspace_billing_history::Column::WorkspaceId.eq(ws.id))
            .filter(workspace_billing_history::Column::Reason.like("ai_usage%"))
            .filter(workspace_billing_history::Column::CreatedAt.gte(month_start))
            .filter(workspace_billing_history::Column::CreatedAt.lt(next_month_start))
            .all(&self.db)
            .await?
            .into_iter()
            .map(|m| m.amount.to_f64().unwrap_or_default())
            .sum::<f64>();
        let month_used = -month_used;

        Ok(WorkspaceBillingCurrentResponse {
            workspace_id: ws.id,
            currency: billing.currency.clone(),
            monthly_quota: billing
                .monthly_quota
                .to_f64()
                .unwrap_or(DEFAULT_MONTHLY_QUOTA),
            balance: billing.balance.to_f64().unwrap_or_default(),
            total_spent: billing.total_spent.to_f64().unwrap_or_default(),
            month_used,
            cycle_start_utc: month_start,
            cycle_end_utc: next_month_start,
            updated_at: billing.updated_at,
            created_at: billing.created_at,
        })
    }

    /// Get workspace billing history.
    pub async fn workspace_billing_history(
        &self,
        ctx: &Session,
        workspace_slug: String,
        query: WorkspaceBillingHistoryQuery,
    ) -> Result<WorkspaceBillingHistoryResponse, AppError> {
        let user_uid = ctx.user().ok_or(AppError::Unauthorized)?;
        let ws = self
            .utils_find_workspace_by_slug(workspace_slug.clone())
            .await?;
        let _ = workspace_membership::Entity::find()
            .filter(workspace_membership::Column::WorkspaceId.eq(ws.id))
            .filter(workspace_membership::Column::UserId.eq(user_uid))
            .filter(workspace_membership::Column::Status.eq("active"))
            .one(&self.db)
            .await?
            .ok_or(AppError::NotWorkspaceMember)?;

        let page = std::cmp::max(query.page.unwrap_or(1), 1);
        let per_page = query.per_page.unwrap_or(20).clamp(1, 200);

        self.ensure_workspace_billing(ws.id, Some(user_uid)).await?;

        let paginator = workspace_billing_history::Entity::find()
            .filter(workspace_billing_history::Column::WorkspaceId.eq(ws.id))
            .order_by_desc(workspace_billing_history::Column::CreatedAt)
            .paginate(&self.db, per_page);
        let total = paginator.num_items().await?;
        let rows = paginator.fetch_page(page - 1).await?;

        let list = rows
            .into_iter()
            .map(|x| WorkspaceBillingHistoryItem {
                uid: x.uid,
                workspace_id: x.workspace_id,
                user_id: x.user_id,
                amount: x.amount.to_f64().unwrap_or_default(),
                currency: x.currency,
                reason: x.reason,
                extra: x.extra.map(|v| v.into()),
                created_at: x.created_at,
            })
            .collect();

        Ok(WorkspaceBillingHistoryResponse {
            page,
            per_page,
            total,
            list,
        })
    }

    /// Add credit to workspace billing (admin action).
    pub async fn workspace_billing_add_credit(
        &self,
        ctx: &Session,
        workspace_slug: String,
        params: WorkspaceBillingAddCreditParams,
    ) -> Result<WorkspaceBillingCurrentResponse, AppError> {
        let user_uid = ctx.user().ok_or(AppError::Unauthorized)?;
        let ws = self
            .utils_find_workspace_by_slug(workspace_slug.clone())
            .await?;
        let _ = workspace_membership::Entity::find()
            .filter(workspace_membership::Column::WorkspaceId.eq(ws.id))
            .filter(workspace_membership::Column::UserId.eq(user_uid))
            .filter(workspace_membership::Column::Status.eq("active"))
            .one(&self.db)
            .await?
            .ok_or(AppError::NotWorkspaceMember)?;

        if params.amount <= 0.0 {
            return Err(AppError::BadRequest("Amount must be positive".to_string()));
        }

        let billing = self.ensure_workspace_billing(ws.id, Some(user_uid)).await?;
        let now_utc = Utc::now();
        let amount_dec =
            Decimal::from_f64_retain(params.amount).unwrap_or(Decimal::ZERO);
        let new_balance = billing.balance + amount_dec;
        let currency = billing.currency.clone();

        let mut updated: workspace_billing::ActiveModel = billing.into();
        updated.balance = Set(new_balance);
        updated.updated_at = Set(now_utc);
        updated.update(&self.db).await?;

        let _ = workspace_billing_history::ActiveModel {
            uid: Set(Uuid::now_v7()),
            workspace_id: Set(ws.id),
            user_id: Set(Some(user_uid)),
            amount: Set(Decimal::from_f64_retain(params.amount).unwrap_or(Decimal::ZERO)),
            currency: Set(currency),
            reason: Set(params.reason.unwrap_or_else(|| "credit_added".to_string())),
            extra: Set(None),
            created_at: Set(now_utc),
        }
        .insert(&self.db)
        .await;

        self.workspace_billing_current(ctx, workspace_slug).await
    }

    /// Ensure workspace billing record exists (create with defaults if not).
    pub async fn ensure_workspace_billing(
        &self,
        workspace_id: Uuid,
        user_uid: Option<Uuid>,
    ) -> Result<workspace_billing::Model, AppError> {
        if let Some(billing) = workspace_billing::Entity::find_by_id(workspace_id)
            .one(&self.db)
            .await?
        {
            return Ok(billing);
        }

        let now_utc = Utc::now();
        // Only first workspace per user gets initial budget ($30)
        let initial_balance = if let Some(uid) = user_uid {
            let existing_workspaces = workspace_membership::Entity::find()
                .filter(workspace_membership::Column::UserId.eq(uid))
                .filter(workspace_membership::Column::Status.eq("active"))
                .all(&self.db)
                .await?;
            if existing_workspaces.len() <= 1 {
                Decimal::from_f64_retain(30.0).unwrap_or(Decimal::ZERO)
            } else {
                Decimal::ZERO
            }
        } else {
            Decimal::ZERO
        };

        let created = workspace_billing::ActiveModel {
            workspace_id: Set(workspace_id),
            balance: Set(initial_balance),
            currency: Set("USD".to_string()),
            monthly_quota: Set(
                Decimal::from_f64_retain(DEFAULT_MONTHLY_QUOTA).unwrap_or(Decimal::ZERO)
            ),
            total_spent: Set(Decimal::ZERO),
            updated_at: Set(now_utc),
            created_at: Set(now_utc),
        };
        Ok(created.insert(&self.db).await?)
    }
}

fn utc_month_bounds(now_utc: DateTime<Utc>) -> Result<(DateTime<Utc>, DateTime<Utc>), AppError> {
    let year = now_utc.year();
    let month = now_utc.month();
    let month_start = NaiveDate::from_ymd_opt(year, month, 1)
        .and_then(|d| d.and_hms_opt(0, 0, 0))
        .map(|d| chrono::TimeZone::from_utc_datetime(&Utc, &d))
        .ok_or_else(|| AppError::InternalServerError("Invalid UTC month start".to_string()))?;
    let (next_year, next_month) = if month == 12 {
        (year + 1, 1)
    } else {
        (year, month + 1)
    };
    let next_month_start = NaiveDate::from_ymd_opt(next_year, next_month, 1)
        .and_then(|d| d.and_hms_opt(0, 0, 0))
        .map(|d| chrono::TimeZone::from_utc_datetime(&Utc, &d))
        .ok_or_else(|| AppError::InternalServerError("Invalid UTC next month start".to_string()))?;
    Ok((month_start, next_month_start))
}