iAs/src/tools/approval.rs

use rand::Rng;
use sha2::{Digest, Sha256};
use sqlx::PgPool;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::{Mutex, oneshot};

/// 审批决策
#[derive(Debug, Clone, PartialEq)]
pub enum ApprovalDecision {
    Approved,
    Rejected,
    Expired,
}

/// 内存中待审批的条目
struct PendingEntry {
    code_hash: String,
    skill_name: String,
    attempts_left: i32,
    /// 通知等待方
    sender: oneshot::Sender<ApprovalDecision>,
}

/// 审批管理器
pub struct ApprovalManager {
    pool: Option<Arc<PgPool>>,
    /// user_id → 该用户的待审批队列
    pending: Arc<Mutex<HashMap<String, Vec<PendingEntry>>>>,
}

impl ApprovalManager {
    pub fn new(pool: Option<Arc<PgPool>>) -> Self {
        Self {
            pool,
            pending: Arc::new(Mutex::new(HashMap::new())),
        }
    }

    /// 创建审批请求
    /// 返回 (确认码明文, 接收器) — 调用方 await 接收器等待用户确认
    pub async fn create(
        &self,
        user_id: &str,
        skill_name: &str,
    ) -> Result<(String, oneshot::Receiver<ApprovalDecision>), String> {
        let code: u32 = rand::thread_rng().gen_range(100000..999999);
        let code_str = code.to_string();
        let code_hash = format!("{:x}", Sha256::digest(code_str.as_bytes()));

        let (tx, rx) = oneshot::channel();

        // 入库
        if let Some(ref pool) = self.pool {
            let id = uuid::Uuid::new_v4();
            let expires_at = chrono::Utc::now() + chrono::Duration::minutes(5);
            let params = serde_json::json!({"name": skill_name});

            let _ = sqlx::query(
                r#"
                INSERT INTO pending_approvals (id, expires_at, user_id, skill_name, params, code_hash, status)
                VALUES ($1, $2, $3, $4, $5, $6, 'pending')
                "#,
            )
            .bind(id)
            .bind(expires_at)
            .bind(user_id)
            .bind(skill_name)
            .bind(&params)
            .bind(&code_hash)
            .execute(pool.as_ref())
            .await;
        }

        // 注册到内存
        let mut map = self.pending.lock().await;
        map.entry(user_id.to_string())
            .or_default()
            .push(PendingEntry {
                code_hash,
                skill_name: skill_name.to_string(),
                attempts_left: 3,
                sender: tx,
            });

        Ok((code_str, rx))
    }

    /// 处理用户回复（由消息循环调用）
    /// 匹配成功时返回 Some(技能名)，匹配失败返回 None
    pub async fn handle_reply(&self, user_id: &str, reply: &str) -> Option<String> {
        let (result, name) = {
            let mut map = self.pending.lock().await;
            let entries = map.get_mut(user_id)?;
            if entries.is_empty() { return None; }

            let reply_hash = format!("{:x}", Sha256::digest(reply.as_bytes()));
            let reply_trimmed = reply.trim();

            if reply_trimmed == "0" || reply_trimmed == "取消" {
                let entry = entries.remove(0);
                let _ = entry.sender.send(ApprovalDecision::Rejected);
                (ApprovalDecision::Rejected, entry.skill_name.clone())
            } else if let Some(idx) = entries.iter().position(|e| e.code_hash == reply_hash) {
                let entry = entries.remove(idx);
                let _ = entry.sender.send(ApprovalDecision::Approved);
                (ApprovalDecision::Approved, entry.skill_name.clone())
            } else if let Some(entry) = entries.first_mut() {
                entry.attempts_left -= 1;
                if entry.attempts_left <= 0 {
                    let entry = entries.remove(0);
                    let n = entry.skill_name.clone();
                    let _ = entry.sender.send(ApprovalDecision::Rejected);
                    (ApprovalDecision::Rejected, n)
                } else {
                    return None;
                }
            } else {
                return None;
            }
        };

        // 更新数据库状态
        if let Some(ref pool) = self.pool {
            let status = match result {
                ApprovalDecision::Approved => "approved",
                _ => "rejected",
            };
            let _ = sqlx::query(
                "UPDATE pending_approvals SET status = $1, consumed_at = NOW() WHERE user_id = $2 AND code_hash = $3 AND status = 'pending'",
            )
            .bind(status)
            .bind(user_id)
            .bind(&format!("{:x}", Sha256::digest(reply.as_bytes())))
            .execute(pool.as_ref())
            .await;
        }

        Some(name)
    }

    /// 清理过期审批（通知等待方 + 更新 DB）
    pub async fn clean_expired(&self) {
        let mut expired = Vec::new();
        {
            let map = self.pending.lock().await;
            for (uid, entries) in map.iter() {
                for (i, e) in entries.iter().enumerate() {
                    if e.sender.is_closed() {
                        expired.push((uid.clone(), i, e.code_hash.clone()));
                    }
                }
            }
        }

        if let Some(ref pool) = self.pool {
            for (_uid, _i, _code_hash) in &expired {
                let _ = sqlx::query(
                    "UPDATE pending_approvals SET status = 'expired' WHERE code_hash = $1 AND status = 'pending'",
                )
                .bind(_code_hash)
                .execute(pool.as_ref())
                .await;
            }
        }

        let mut map = self.pending.lock().await;
        for (uid, i, _) in expired.iter().rev() {
            if let Some(entries) = map.get_mut(uid) {
                if *i < entries.len() {
                    let entry = entries.remove(*i);
                    let _ = entry.sender.send(ApprovalDecision::Expired);
                }
            }
        }
        map.retain(|_, v| !v.is_empty());
    }

    /// 获取某个用户的待审批条数
    pub async fn pending_count(&self, user_id: &str) -> usize {
        let map = self.pending.lock().await;
        map.get(user_id).map(|v| v.len()).unwrap_or(0)
    }
}