hms/crates/erp-ai/src/service/quota.rs

use sea_orm::{ColumnTrait, EntityTrait, FromQueryResult, QueryFilter, Statement};
use uuid::Uuid;

use crate::entity::ai_tenant_config;
use crate::error::{AiError, AiResult};

pub struct QuotaService {
    db: sea_orm::DatabaseConnection,
    enabled: bool,
}

impl QuotaService {
    pub fn new(db: sea_orm::DatabaseConnection, enabled: bool) -> Self {
        Self { db, enabled }
    }

    pub async fn get_tenant_config(
        &self,
        tenant_id: Uuid,
    ) -> AiResult<Option<ai_tenant_config::Model>> {
        let config = ai_tenant_config::Entity::find()
            .filter(ai_tenant_config::Column::TenantId.eq(tenant_id))
            .filter(ai_tenant_config::Column::DeletedAt.is_null())
            .filter(ai_tenant_config::Column::IsEnabled.eq(true))
            .one(&self.db)
            .await?;
        Ok(config)
    }

    pub async fn check_quota(&self, tenant_id: Uuid, patient_id: Option<Uuid>) -> AiResult<()> {
        if !self.enabled {
            return Ok(());
        }

        let config = match self.get_tenant_config(tenant_id).await? {
            Some(c) => c,
            None => return Ok(()),
        };

        let month_usage = self.get_monthly_token_usage(tenant_id).await?;
        if month_usage >= config.monthly_token_budget {
            return Err(AiError::QuotaExhausted {
                tenant_id,
                reason: format!(
                    "月度 Token 预算已耗尽 ({}/{})",
                    month_usage, config.monthly_token_budget
                ),
            });
        }

        if let Some(pid) = patient_id {
            let daily_count = self.get_daily_patient_count(tenant_id, pid).await?;
            if daily_count >= config.daily_patient_limit as i64 {
                return Err(AiError::QuotaExhausted {
                    tenant_id,
                    reason: format!(
                        "患者 {} 今日分析次数已达上限 ({}/{})",
                        pid, daily_count, config.daily_patient_limit
                    ),
                });
            }
        }

        Ok(())
    }

    async fn get_monthly_token_usage(&self, tenant_id: Uuid) -> AiResult<i64> {
        #[derive(Debug, FromQueryResult)]
        struct TokenSum {
            total_tokens: i64,
        }

        let sql = r#"
            SELECT COALESCE(SUM(input_tokens + output_tokens), 0) AS total_tokens
            FROM ai_usage
            WHERE tenant_id = $1
              AND created_at >= date_trunc('month', CURRENT_DATE)
        "#;

        let result: Option<TokenSum> = TokenSum::find_by_statement(Statement::from_sql_and_values(
            sea_orm::DatabaseBackend::Postgres,
            sql,
            [tenant_id.into()],
        ))
        .one(&self.db)
        .await?;

        Ok(result.map(|r| r.total_tokens).unwrap_or(0))
    }

    async fn get_daily_patient_count(&self, tenant_id: Uuid, patient_id: Uuid) -> AiResult<i64> {
        #[derive(Debug, FromQueryResult)]
        struct CountResult {
            count: i64,
        }

        let sql = r#"
            SELECT COUNT(*) AS count
            FROM ai_analysis_results
            WHERE tenant_id = $1
              AND patient_id = $2
              AND deleted_at IS NULL
              AND created_at >= CURRENT_DATE
        "#;

        let result: Option<CountResult> =
            CountResult::find_by_statement(Statement::from_sql_and_values(
                sea_orm::DatabaseBackend::Postgres,
                sql,
                [tenant_id.into(), patient_id.into()],
            ))
            .one(&self.db)
            .await?;

        Ok(result.map(|r| r.count).unwrap_or(0))
    }

    pub async fn get_usage_summary(&self, tenant_id: Uuid) -> AiResult<QuotaSummary> {
        let config = self.get_tenant_config(tenant_id).await.unwrap_or(None);
        let budget = config
            .as_ref()
            .map(|c| c.monthly_token_budget)
            .unwrap_or(1_000_000);

        let used = self.get_monthly_token_usage(tenant_id).await.unwrap_or(0);

        Ok(QuotaSummary {
            tenant_id,
            monthly_budget: budget,
            monthly_used: used,
            daily_patient_limit: config.as_ref().map(|c| c.daily_patient_limit).unwrap_or(50),
        })
    }
}

#[derive(Debug, Clone, serde::Serialize)]
pub struct QuotaSummary {
    pub tenant_id: Uuid,
    pub monthly_budget: i64,
    pub monthly_used: i64,
    pub daily_patient_limit: i32,
}

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

    #[test]
    fn quota_summary_serialization() {
        let summary = QuotaSummary {
            tenant_id: Uuid::now_v7(),
            monthly_budget: 1_000_000,
            monthly_used: 250_000,
            daily_patient_limit: 50,
        };
        let json = serde_json::to_value(&summary).unwrap();
        assert_eq!(json["monthly_budget"], 1_000_000);
        assert_eq!(json["monthly_used"], 250_000);
    }

    #[test]
    fn quota_error_contains_reason() {
        let tid = Uuid::now_v7();
        let err = AiError::QuotaExhausted {
            tenant_id: tid,
            reason: "月度预算耗尽".into(),
        };
        let msg = err.to_string();
        assert!(msg.contains("月度预算耗尽"));
    }
}