gitdataai/libs/service/project/billing.rs

use crate::AppService;
use crate::error::AppError;
use chrono::{DateTime, Datelike, NaiveDate, Utc};
use models::Decimal;
use models::projects::{project_billing, project_billing_history};
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;

const DEFAULT_PROJECT_MONTHLY_CREDIT: f64 = 10.0;

#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
pub struct ProjectBillingCurrentResponse {
    pub project_uid: Uuid,
    pub currency: String,
    pub monthly_quota: f64,
    pub balance: 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 ProjectBillingHistoryQuery {
    pub page: Option<u64>,
    pub per_page: Option<u64>,
}

#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
pub struct ProjectBillingHistoryItem {
    pub uid: Uuid,
    pub project_uid: Uuid,
    pub user_uid: 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 ProjectBillingHistoryResponse {
    pub page: u64,
    pub per_page: u64,
    pub total: u64,
    pub list: Vec<ProjectBillingHistoryItem>,
}

impl AppService {
    pub async fn project_billing_current(
        &self,
        ctx: &Session,
        project_name: String,
    ) -> Result<ProjectBillingCurrentResponse, AppError> {
        let user_uid = ctx.user().ok_or(AppError::Unauthorized)?;
        let project = self.utils_find_project_by_name(project_name).await?;
        self.check_project_access(project.id, user_uid).await?;

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

        let month_used = project_billing_history::Entity::find()
            .filter(project_billing_history::Column::Project.eq(project.id))
            .filter(project_billing_history::Column::Reason.like("ai_usage%"))
            .filter(project_billing_history::Column::CreatedAt.gte(month_start))
            .filter(project_billing_history::Column::CreatedAt.lt(next_month_start))
            .all(&self.db)
            .await?
            .into_iter()
            .map(|m| m.amount)
            .sum::<Decimal>();
        let month_used = -month_used;

        Ok(ProjectBillingCurrentResponse {
            project_uid: project.id,
            currency: billing.currency,
            monthly_quota: DEFAULT_PROJECT_MONTHLY_CREDIT,
            balance: billing.balance.to_f64().unwrap_or_default(),
            month_used: month_used.to_f64().unwrap_or_default(),
            cycle_start_utc: month_start,
            cycle_end_utc: next_month_start,
            updated_at: billing.updated_at,
            created_at: billing.created_at,
        })
    }

    pub async fn project_billing_history(
        &self,
        ctx: &Session,
        project_name: String,
        query: ProjectBillingHistoryQuery,
    ) -> Result<ProjectBillingHistoryResponse, AppError> {
        let user_uid = ctx.user().ok_or(AppError::Unauthorized)?;
        let project = self.utils_find_project_by_name(project_name).await?;
        self.check_project_access(project.id, user_uid).await?;

        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_project_billing(project.id, Some(user_uid))
            .await?;

        let paginator = project_billing_history::Entity::find()
            .filter(project_billing_history::Column::Project.eq(project.id))
            .order_by_desc(project_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| ProjectBillingHistoryItem {
                uid: x.uid,
                project_uid: x.project,
                user_uid: x.user,
                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(ProjectBillingHistoryResponse {
            page,
            per_page,
            total,
            list,
        })
    }

    pub async fn ensure_project_billing(
        &self,
        project_uid: Uuid,
        user_uid: Option<Uuid>,
    ) -> Result<project_billing::Model, AppError> {
        if let Some(billing) = project_billing::Entity::find_by_id(project_uid)
            .one(&self.db)
            .await?
        {
            return Ok(billing);
        }

        let now_utc = Utc::now();
        // Only first project per user gets initial budget ($10)
        let initial_balance = if let Some(uid) = user_uid {
            let existing_projects = models::projects::project::Entity::find()
                .filter(models::projects::project::Column::CreatedBy.eq(uid))
                .all(&self.db)
                .await?;
            if existing_projects.len() <= 1 {
                Decimal::from_f64_retain(DEFAULT_PROJECT_MONTHLY_CREDIT).unwrap_or(Decimal::ZERO)
            } else {
                Decimal::ZERO
            }
        } else {
            Decimal::ZERO
        };

        let created = project_billing::ActiveModel {
            project: Set(project_uid),
            balance: Set(initial_balance),
            currency: Set("USD".to_string()),
            user: Set(user_uid),
            updated_at: Set(now_utc),
            created_at: Set(now_utc),
            ..Default::default()
        };
        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))
}