limiteron 0.3.0-rc.4

// Copyright (c) 2026 Kirky.X
// SPDX-License-Identifier: MIT
//! DBNexusQuotaStorageAdapter - DBNexus-based implementation of QuotaStorage trait
//!
//! This adapter provides a complete QuotaStorage trait implementation using DBNexus
//! for all quota management operations.

use crate::dbnexus_entities::{
    QuotaColumn, QuotaRecordActiveModel, QuotaRecordModel, create_quota_key,
};
use crate::error::{ConsumeResult, StorageError};
use crate::storage::{QuotaInfo, QuotaStorage};
use async_trait::async_trait;
use chrono::{Duration as ChronoDuration, Utc};
use dbnexus::{Condition, DbPool, Session};
use sea_orm::entity::prelude::*;
use sea_orm::{DatabaseBackend, Statement};
use std::sync::Arc;
use std::time::Duration as StdDuration;

/// Read the `RETURNING consumed` column from an atomic quota update
fn row_consumed(row: &sea_orm::QueryResult) -> Result<u64, StorageError> {
    row.try_get::<i64>("", "consumed")
        .map(|v| v as u64)
        .map_err(|e| StorageError::QueryError(format!("Failed to read consumed: {}", e)))
}

/// DBNexus-based quota storage adapter
pub struct DBNexusQuotaStorageAdapter {
    pool: Arc<DbPool>,
}

impl DBNexusQuotaStorageAdapter {
    /// Create a new DBNexusQuotaStorageAdapter
    pub fn new(pool: Arc<DbPool>) -> Self {
        Self { pool }
    }

    /// Get a session from the pool
    async fn get_session(&self) -> Result<Session, StorageError> {
        self.pool
            .get_session("admin")
            .await
            .map_err(|e| StorageError::ConnectionError(e.to_string()))
    }

    /// Get a connection from the session (for direct sea-orm operations)
    fn get_conn(session: &Session) -> Result<&DatabaseConnection, StorageError> {
        session
            .connection()
            .map_err(|e| StorageError::ConnectionError(e.to_string()))
    }

    /// Convert model to QuotaInfo
    fn model_to_info(model: &QuotaRecordModel) -> QuotaInfo {
        QuotaInfo {
            consumed: model.consumed as u64,
            limit: model.limit as u64,
            window_start: model.window_start,
            window_end: model.window_end,
        }
    }

    /// Map dbnexus DbError to StorageError
    fn map_err(e: dbnexus::DbError) -> StorageError {
        StorageError::QueryError(e.to_string())
    }

    /// Execute a statement that returns at most one row (None if no match)
    async fn query_optional(
        conn: &DatabaseConnection,
        sql: &str,
        values: impl IntoIterator<Item = sea_orm::Value>,
    ) -> Result<Option<sea_orm::QueryResult>, StorageError> {
        conn.query_one_raw(Statement::from_sql_and_values(
            DatabaseBackend::Postgres,
            sql,
            values,
        ))
        .await
        .map_err(|e| {
            StorageError::QueryError(format!("Failed to execute atomic quota update: {}", e))
        })
    }

    /// Build an allowed ConsumeResult from the post-consumption ledger value
    /// （diting 简化：改用共享的 `ConsumeResult::allowed` 构造器）
    fn allowed_result(new_consumed: u64, limit: u64) -> ConsumeResult {
        ConsumeResult::allowed(new_consumed, limit)
    }
}

#[async_trait]
impl QuotaStorage for DBNexusQuotaStorageAdapter {
    /// Get quota info for a user and resource
    async fn get_quota(
        &self,
        user_id: &str,
        resource: &str,
    ) -> Result<Option<QuotaInfo>, StorageError> {
        let session = self.get_session().await?;
        let quota_key = create_quota_key(user_id, resource);
        let now = Utc::now();

        let condition = Condition::all()
            .add(QuotaColumn::QuotaKey.eq(quota_key))
            .add(QuotaColumn::WindowEnd.gt(now));

        let records = QuotaRecordModel::find_by_condition(&session, condition)
            .await
            .map_err(Self::map_err)?;

        Ok(records.into_iter().next().map(|m| Self::model_to_info(&m)))
    }

    /// Consume quota
    ///
    /// 原子条件 UPDATE（A3）：限额检查与累加在单条 SQL 内由数据库完成
    /// （`WHERE consumed + cost <= "limit"` + `RETURNING`），消除
    /// read-check-write 竞态下的静默超额放行。未命中时区分「超限拒绝」
    /// 与「无活跃记录」；后者经新窗口原子重启（复用过期行）或守卫式
    /// 插入处理，两条路径都以 `ON CONFLICT` 兜底并发首触（diting Medium：
    /// 旧实现并发 INSERT 撞 quota_key UNIQUE 会向请求返回错误）。
    async fn consume(
        &self,
        user_id: &str,
        resource: &str,
        cost: u64,
        limit: u64,
        window: StdDuration,
    ) -> Result<ConsumeResult, StorageError> {
        let session = self.get_session().await?;
        let conn = Self::get_conn(&session)?;
        let quota_key = create_quota_key(user_id, resource);
        let now = Utc::now();
        let chrono_window =
            ChronoDuration::from_std(window).unwrap_or_else(|_| ChronoDuration::days(365));
        let window_end = now + chrono_window;

        // 活跃窗口原子累加：命中即允许并返回累加后的 consumed。
        // 限额取 LEAST(存储列, 调用方参数)（diting Low）：运行期调低
        // 配额上限时，活跃窗口也按新上限裁决。
        const CONSUME_SQL: &str = r#"
            UPDATE limiteron_quotas
            SET consumed = consumed + $2, updated_at = $4
            WHERE quota_key = $1 AND window_end > $4 AND consumed + $2 <= LEAST("limit", $3)
            RETURNING consumed
        "#;
        // 守卫式插入：仅当冲突行确实过期时以全新窗口覆盖（并发首触时
        // 输家不再撞 UNIQUE 报错，而是回到循环顶部走原子累加路径）
        const INSERT_SQL: &str = r#"
            INSERT INTO limiteron_quotas
                (user_id, resource, quota_key, "limit", consumed,
                 window_start, window_end, created_at, updated_at)
            VALUES ($1, $2, $3, $4, $5, $6, $7, $6, $6)
            ON CONFLICT (quota_key) DO UPDATE SET
                consumed = EXCLUDED.consumed,
                "limit" = EXCLUDED."limit",
                window_start = EXCLUDED.window_start,
                window_end = EXCLUDED.window_end,
                updated_at = EXCLUDED.updated_at
            WHERE limiteron_quotas.window_end <= EXCLUDED.window_start
            RETURNING consumed
        "#;
        // 新窗口原子重启：同 key 的过期行在「窗口确实过期 + 不超限」
        // 条件下原子复用（保留行身份与创建时间）
        const RESTART_SQL: &str = r#"
            UPDATE limiteron_quotas
            SET consumed = $2, "limit" = $3, window_start = $4, window_end = $5, updated_at = $4
            WHERE quota_key = $1 AND window_end <= $4 AND $2 <= $3
            RETURNING consumed
        "#;

        // 有界重试：并发首触时败方的 INSERT 冲突由守卫吸收（0 行返回），
        // 回到顶部重跑原子累加即可命中胜方建立的活跃行
        for _ in 0..3 {
            if let Some(row) = Self::query_optional(
                conn,
                CONSUME_SQL,
                [
                    quota_key.clone().into(),
                    (cost as i64).into(),
                    (limit as i64).into(),
                    now.into(),
                ],
            )
            .await?
            {
                return Ok(Self::allowed_result(row_consumed(&row)?, limit));
            }

            // 未命中：区分超限拒绝与无活跃记录（新窗口/首次消费）
            let condition = Condition::all()
                .add(QuotaColumn::QuotaKey.eq(quota_key.clone()))
                .add(QuotaColumn::WindowEnd.gt(now));

            if let Some(record) = QuotaRecordModel::find_by_condition(&session, condition)
                .await
                .map_err(Self::map_err)?
                .into_iter()
                .next()
            {
                // 超限拒绝：给出基于当前账本的准确 remaining/usage
                let consumed = record.consumed as u64;
                return Ok(ConsumeResult::rejected(consumed, limit));
            }

            if cost > limit {
                return Ok(ConsumeResult::rejected(0, limit));
            }

            // 先试新窗口原子重启（复用过期行，保留行身份与创建时间）
            if let Some(row) = Self::query_optional(
                conn,
                RESTART_SQL,
                [
                    quota_key.clone().into(),
                    (cost as i64).into(),
                    (limit as i64).into(),
                    now.into(),
                    window_end.into(),
                ],
            )
            .await?
            {
                let _ = row_consumed(&row)?;
                return Ok(Self::allowed_result(cost, limit));
            }

            // 全新 key：守卫式插入。并发首触时败方 INSERT 撞 quota_key
            // UNIQUE，但 ON CONFLICT 的 WHERE 过期守卫使其成为 0 行更新
            // （不再向请求抛 UNIQUE 错误），回到循环顶部走原子累加路径。
            if let Some(row) = Self::query_optional(
                conn,
                INSERT_SQL,
                [
                    user_id.into(),
                    resource.into(),
                    quota_key.clone().into(),
                    (limit as i64).into(),
                    (cost as i64).into(),
                    now.into(),
                    window_end.into(),
                ],
            )
            .await?
            {
                let _ = row_consumed(&row)?;
                return Ok(Self::allowed_result(cost, limit));
            }
            // 0 行返回 = 冲突行仍活跃（竞态），循环重试
        }

        Err(StorageError::QueryError(
            "并发配额初始化冲突重试耗尽，请重试请求".to_string(),
        ))
    }

    /// Reset quota
    async fn reset(
        &self,
        user_id: &str,
        resource: &str,
        limit: u64,
        window: StdDuration,
    ) -> Result<(), StorageError> {
        let session = self.get_session().await?;
        let conn = Self::get_conn(&session)?;
        let quota_key = create_quota_key(user_id, resource);
        let now = Utc::now();
        let chrono_window =
            ChronoDuration::from_std(window).unwrap_or_else(|_| ChronoDuration::days(365));
        let window_end = now + chrono_window;

        let model = QuotaRecordModel {
            id: 0,
            user_id: user_id.to_string(),
            resource: resource.to_string(),
            quota_key: quota_key.clone(),
            limit: limit as i64,
            consumed: 0,
            window_start: now,
            window_end,
            created_at: now,
            updated_at: now,
        };

        // 先查现有记录：存在则重置并更新（同一 quota_key 直接 insert 会 duplicate 崩溃），
        // 不存在才插入新记录（主键由 BIGSERIAL 序列生成）。
        let condition = Condition::all().add(QuotaColumn::QuotaKey.eq(quota_key));
        let existing = QuotaRecordModel::find_by_condition(&session, condition)
            .await
            .map_err(Self::map_err)?
            .into_iter()
            .next();

        if let Some(record) = existing {
            let mut am: QuotaRecordActiveModel = record.into();
            am.consumed = sea_orm::Set(0);
            am.limit = sea_orm::Set(limit as i64);
            am.window_start = sea_orm::Set(now);
            am.window_end = sea_orm::Set(window_end);
            am.created_at = sea_orm::Set(now);
            am.updated_at = sea_orm::Set(now);
            am.save(conn)
                .await
                .map_err(|e| StorageError::QueryError(format!("Failed to reset quota: {}", e)))?;
            return Ok(());
        }

        let mut active_model: QuotaRecordActiveModel = model.into();
        // 主键由 BIGSERIAL 序列生成：显式 Set(0) 会导致重复主键崩溃
        active_model.id = sea_orm::NotSet;
        active_model
            .insert(conn)
            .await
            .map_err(|e| StorageError::QueryError(format!("Failed to reset quota: {}", e)))?;

        Ok(())
    }
}

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

    fn make_model(user_id: &str, resource: &str, limit: u64, consumed: u64) -> QuotaRecordModel {
        let window_start = Utc.with_ymd_and_hms(2026, 1, 1, 0, 0, 0).unwrap();
        let window_end = window_start + ChronoDuration::hours(1);
        QuotaRecordModel {
            id: 1,
            user_id: user_id.to_string(),
            resource: resource.to_string(),
            quota_key: create_quota_key(user_id, resource),
            limit: limit as i64,
            consumed: consumed as i64,
            window_start,
            window_end,
            created_at: window_start,
            updated_at: window_start,
        }
    }

    #[test]
    fn test_model_to_info_basic() {
        let model = make_model("user1", "api_requests", 1000, 250);
        let info = DBNexusQuotaStorageAdapter::model_to_info(&model);
        assert_eq!(info.consumed, 250);
        assert_eq!(info.limit, 1000);
        assert_eq!(info.window_start, model.window_start);
        assert_eq!(info.window_end, model.window_end);
    }

    #[test]
    fn test_model_to_info_zero_consumed() {
        let model = make_model("user2", "storage_mb", 500, 0);
        let info = DBNexusQuotaStorageAdapter::model_to_info(&model);
        assert_eq!(info.consumed, 0);
        assert_eq!(info.limit, 500);
    }

    #[test]
    fn test_model_to_info_full_consumed() {
        let model = make_model("user3", "api_calls", 100, 100);
        let info = DBNexusQuotaStorageAdapter::model_to_info(&model);
        assert_eq!(info.consumed, 100);
        assert_eq!(info.limit, 100);
    }

    #[test]
    fn test_model_to_info_preserves_window() {
        let model = make_model("user4", "resource", 50, 25);
        let info = DBNexusQuotaStorageAdapter::model_to_info(&model);
        assert_eq!(info.window_start, model.window_start);
        assert_eq!(info.window_end, model.window_end);
    }

    #[test]
    fn test_map_err_db_error_to_storage_error() {
        let db_err = dbnexus::DbError::new(sea_orm::DbErr::Custom("quota query error".to_string()));
        let storage_err = DBNexusQuotaStorageAdapter::map_err(db_err);
        match storage_err {
            StorageError::QueryError(msg) => {
                assert!(msg.contains("quota query error"));
            }
            other => panic!("expected QueryError, got {:?}", other),
        }
    }
}