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

use sea_orm::{
    ActiveModelTrait, ColumnTrait, EntityTrait, FromQueryResult, QueryFilter, QuerySelect, Set,
};
use uuid::Uuid;

use crate::entity::ai_analysis;
use crate::entity::ai_usage;
use crate::entity::ai_usage_daily;
use crate::error::AiResult;

pub struct UsageService {
    pub db: sea_orm::DatabaseConnection,
}

impl UsageService {
    pub fn new(db: sea_orm::DatabaseConnection) -> Self {
        Self { db }
    }

    #[allow(clippy::too_many_arguments)]
    pub async fn log_usage(
        &self,
        tenant_id: Uuid,
        provider: &str,
        model: &str,
        analysis_type: &str,
        input_tokens: u32,
        output_tokens: u32,
        duration_ms: u64,
        cost_cents: i32,
        is_cache_hit: bool,
    ) -> AiResult<ai_usage::Model> {
        let id = Uuid::now_v7();
        let active = ai_usage::ActiveModel {
            id: Set(id),
            tenant_id: Set(tenant_id),
            provider: Set(provider.into()),
            model: Set(model.into()),
            analysis_type: Set(analysis_type.into()),
            input_tokens: Set(input_tokens as i32),
            output_tokens: Set(output_tokens as i32),
            duration_ms: Set(duration_ms as i32),
            cost_cents: Set(cost_cents),
            is_cache_hit: Set(is_cache_hit),
            created_at: Set(chrono::Utc::now()),
        };
        Ok(active.insert(&self.db).await?)
    }

    /// 用量概览
    pub async fn get_overview(&self, tenant_id: Uuid) -> AiResult<UsageOverview> {
        let result = ai_analysis::Entity::find()
            .filter(ai_analysis::Column::TenantId.eq(tenant_id))
            .filter(ai_analysis::Column::Status.eq("completed"))
            .filter(ai_analysis::Column::DeletedAt.is_null())
            .select_only()
            .column_as(ai_analysis::Column::Id.count(), "total_count")
            .into_model::<UsageOverview>()
            .one(&self.db)
            .await?
            .unwrap_or(UsageOverview { total_count: 0 });
        Ok(result)
    }

    /// 按分析类型统计
    pub async fn get_by_type(&self, tenant_id: Uuid) -> AiResult<Vec<TypeCount>> {
        let result = ai_analysis::Entity::find()
            .filter(ai_analysis::Column::TenantId.eq(tenant_id))
            .filter(ai_analysis::Column::Status.eq("completed"))
            .filter(ai_analysis::Column::DeletedAt.is_null())
            .select_only()
            .column(ai_analysis::Column::AnalysisType)
            .column_as(ai_analysis::Column::Id.count(), "count")
            .group_by(ai_analysis::Column::AnalysisType)
            .into_model::<TypeCount>()
            .all(&self.db)
            .await?;
        Ok(result)
    }

    /// 按日期范围查询日聚合用量
    pub async fn get_daily_usage(
        &self,
        tenant_id: Uuid,
        start_date: chrono::NaiveDate,
        end_date: chrono::NaiveDate,
    ) -> AiResult<Vec<DailyUsageRow>> {
        let rows = ai_usage_daily::Entity::find()
            .filter(ai_usage_daily::Column::TenantId.eq(tenant_id))
            .filter(ai_usage_daily::Column::Date.gte(start_date))
            .filter(ai_usage_daily::Column::Date.lte(end_date))
            .all(&self.db)
            .await?;
        Ok(rows
            .into_iter()
            .map(|r| DailyUsageRow {
                date: r.date,
                feature: r.feature,
                provider: r.provider,
                model: r.model,
                total_calls: r.total_calls,
                total_input_tokens: r.total_input_tokens,
                total_output_tokens: r.total_output_tokens,
                total_cost_cents: r.total_cost_cents,
            })
            .collect())
    }

    /// 聚合指定日期的用量到日聚合表（由定时任务调用）
    pub async fn aggregate_daily(&self, tenant_id: Uuid, date: chrono::NaiveDate) -> AiResult<()> {
        let date_start = date.and_hms_opt(0, 0, 0).unwrap_or_default();
        let date_end = date_start + chrono::Duration::days(1);

        // 从 ai_usage 按分析类型聚合
        #[derive(Debug, FromQueryResult)]
        struct AggRow {
            analysis_type: String,
            total_calls: i64,
            total_input_tokens: i64,
            total_output_tokens: i64,
            total_cost_cents: i64,
        }

        let rows: Vec<AggRow> = ai_usage::Entity::find()
            .filter(ai_usage::Column::TenantId.eq(tenant_id))
            .filter(ai_usage::Column::CreatedAt.gte(date_start))
            .filter(ai_usage::Column::CreatedAt.lt(date_end))
            .select_only()
            .column(ai_usage::Column::AnalysisType)
            .column_as(ai_usage::Column::Id.count(), "total_calls")
            .column_as(ai_usage::Column::InputTokens.sum(), "total_input_tokens")
            .column_as(ai_usage::Column::OutputTokens.sum(), "total_output_tokens")
            .column_as(ai_usage::Column::CostCents.sum(), "total_cost_cents")
            .group_by(ai_usage::Column::AnalysisType)
            .into_model::<AggRow>()
            .all(&self.db)
            .await?;

        for row in &rows {
            let id = Uuid::now_v7();
            let active = ai_usage_daily::ActiveModel {
                id: Set(id),
                tenant_id: Set(tenant_id),
                date: Set(date),
                feature: Set(row.analysis_type.clone()),
                provider: Set("aggregated".into()),
                model: Set("mixed".into()),
                total_calls: Set(row.total_calls as i32),
                total_input_tokens: Set(row.total_input_tokens),
                total_output_tokens: Set(row.total_output_tokens),
                total_cost_cents: Set(row.total_cost_cents),
                created_at: Set(chrono::Utc::now()),
            };
            active.insert(&self.db).await?;
        }

        Ok(())
    }
}

#[derive(Debug, FromQueryResult)]
pub struct UsageOverview {
    pub total_count: i64,
}

#[derive(Debug, FromQueryResult)]
pub struct TypeCount {
    pub analysis_type: String,
    pub count: i64,
}

#[derive(Debug, Clone, serde::Serialize)]
pub struct DailyUsageRow {
    pub date: chrono::NaiveDate,
    pub feature: String,
    pub provider: String,
    pub model: String,
    pub total_calls: i32,
    pub total_input_tokens: i64,
    pub total_output_tokens: i64,
    pub total_cost_cents: i64,
}