sz-orm-queue 1.2.2

//! Redis 真实队列 provider 实现（基于 `redis` crate 0.27，启用 `tokio-comp` + `streams`）
//!
//! 通过 [`RedisMode`] 支持三种工作模式：
//! - **List**：基于 List 的点对点队列（`RPUSH` + `BLPOP`），FIFO，消费即出队，无 ACK
//! - **PubSub**：基于 Pub/Sub 的广播（`PUBLISH` + `SUBSCRIBE`），无持久化，无 ACK
//! - **Stream**：基于 Stream 的可靠队列（`XADD` + `XREADGROUP` + `XACK`），消费者组，支持 ACK
//!
//! 连接复用：List / Stream 模式共用一条 `MultiplexedConnection`（内部基于 Arc，可低成本克隆）；
//! PubSub 模式按 Redis 协议要求为每个 topic 另起一条独立订阅连接。
//!
//! 依赖：`redis` crate 启用 `tokio-comp` + `streams` feature。

use crate::error::MqError;
use crate::queue::{Message, MessageQueue, RedisConfig, RedisMode};
use async_trait::async_trait;
use futures::StreamExt;
use redis::aio::{MultiplexedConnection, PubSubStream};
use redis::{AsyncCommands, RedisError};
use std::collections::HashMap;
use std::pin::Pin;
use std::sync::Arc;
use tokio::sync::Mutex;

/// Stream 模式消费者组的默认起始 ID：`"0"` 表示从最早未消费消息开始消费
const STREAM_START_FROM_BEGINNING: &str = "0";
/// List 模式 `BLPOP` 的阻塞超时秒数（避免永久阻塞 trait 调用）
const LIST_BLOCK_TIMEOUT_SECS: f64 = 1.0;
/// PubSub / Stream 模式短轮询超时（毫秒）：无消息时及时返回 `None`
const POLL_TIMEOUT_MS: u64 = 100;

/// Redis 真实队列 provider
///
/// 通过 [`RedisQueueProvider::connect`] 传入完整 [`RedisConfig`] 建立连接后即可使用
/// [`MessageQueue`] trait 方法；也可用 [`RedisQueueProvider::new`] 仅凭 URL 快速接入
/// （默认 List 模式）。
pub struct RedisQueueProvider {
    /// 队列配置（工作模式、消费者组等）
    config: RedisConfig,
    /// 多路复用连接（List / Stream 模式共用，`MultiplexedConnection` 内部已基于 Arc，
    /// 可低成本克隆，因此无需再套 `Arc`；redis 的 `AsyncCommands` 方法需要 `&mut self`）
    conn: MultiplexedConnection,
    /// PubSub 模式：每个 topic 一条订阅流（独立连接 `into_on_message` 得到的 owned stream）
    pubsub_streams: Arc<Mutex<HashMap<String, Pin<Box<PubSubStream>>>>>,
    /// Stream 模式：in-flight 消息元信息（message_id -> (stream_key, group)），供 ack 使用
    in_flight: Arc<Mutex<HashMap<String, (String, String)>>>,
}
impl RedisQueueProvider {
    /// 便捷构造：仅凭 URL 连接 Redis（使用默认 List 模式）
    ///
    /// 等价于 `connect(RedisConfig { url: Some(url.into()), ..Default::default() })`。
    /// 如需 Stream / PubSub 模式或自定义消费者组，请使用 [`connect`](Self::connect)。
    pub async fn new(url: impl Into<String>) -> Result<Self, MqError> {
        Self::connect(RedisConfig {
            url: Some(url.into()),
            ..RedisConfig::default()
        })
        .await
    }

    /// 建立到 Redis 的连接并构造队列实例
    ///
    /// - 优先使用 `config.url`；未设置时默认 `redis://127.0.0.1:6379/0`
    /// - Stream 模式下未配置消费者组/消费者名时给出默认值
    pub async fn connect(mut config: RedisConfig) -> Result<Self, MqError> {
        let url = config
            .url
            .clone()
            .unwrap_or_else(|| "redis://127.0.0.1:6379/0".to_string());
        let client = redis::Client::open(url.as_str())
            .map_err(|e| MqError::Connection(format!("Redis client open failed: {e}")))?;
        let conn = client
            .get_multiplexed_async_connection()
            .await
            .map_err(|e| MqError::Connection(format!("Redis connect failed: {e}")))?;
        // Stream 模式需要消费者组/消费者名；未配置时补默认值
        if config.consumer_group.is_none() {
            config.consumer_group = Some("sz-orm-queue-group".to_string());
        }
        if config.consumer_name.is_none() {
            config.consumer_name = Some("consumer-1".to_string());
        }
        Ok(Self {
            config,
            conn,
            pubsub_streams: Arc::new(Mutex::new(HashMap::new())),
            in_flight: Arc::new(Mutex::new(HashMap::new())),
        })
    }

    /// 返回当前工作模式
    pub fn mode(&self) -> RedisMode {
        self.config.mode
    }

    /// 返回消费者组名称（Stream 模式使用）
    fn group(&self) -> &str {
        self.config
            .consumer_group
            .as_deref()
            .unwrap_or("sz-orm-queue-group")
    }

    /// 返回消费者名称（Stream 模式使用）
    fn consumer(&self) -> &str {
        self.config
            .consumer_name
            .as_deref()
            .unwrap_or("consumer-1")
    }

    // ---------------- List 模式 ----------------

    /// List 模式发布：`RPUSH` 将消息追加到列表尾部，配合 `BLPOP`（弹出头部）实现 FIFO
    async fn publish_list(&self, topic: &str, message: &[u8]) -> Result<(), MqError> {
        let mut conn = self.conn.clone();
        let _: i64 = conn
            .rpush(topic, message.to_vec())
            .await
            .map_err(|e| MqError::Publish(format!("Redis RPUSH failed: {e}")))?;
        Ok(())
    }

    /// List 模式消费：`BLPOP` 阻塞弹出列表头部（超时 `LIST_BLOCK_TIMEOUT_SECS` 秒）
    ///
    /// 返回 `None` 表示当前无消息。List 模式消费即出队，无需 ack。
    async fn consume_list(&self, topic: &str) -> Result<Option<Message>, MqError> {
        let mut conn = self.conn.clone();
        let result: Option<(String, Vec<u8>)> = conn
            .blpop(topic, LIST_BLOCK_TIMEOUT_SECS)
            .await
            .map_err(|e| MqError::Connection(format!("Redis BLPOP failed: {e}")))?;
        Ok(result.map(|(_, payload)| Message {
            topic: topic.to_string(),
            payload,
            key: None,
            timestamp: current_timestamp_millis(),
            headers: HashMap::new(),
            // List 模式无 entry id；用临时 id，ack 为 no-op
            id: format!("redis-list-{}", current_timestamp_millis()),
            retry_count: 0,
        }))
    }
    // ---------------- PubSub 模式 ----------------

    /// PubSub 模式发布：`PUBLISH` 将消息推送到所有订阅者（无持久化）
    async fn publish_pubsub(&self, topic: &str, message: &[u8]) -> Result<(), MqError> {
        let mut conn = self.conn.clone();
        let _: i64 = conn
            .publish(topic, message.to_vec())
            .await
            .map_err(|e| MqError::Publish(format!("Redis PUBLISH failed: {e}")))?;
        Ok(())
    }

    /// PubSub 模式消费：从已订阅的 `PubSubStream` 中拉取一条消息（短超时，无消息返回 `None`）
    ///
    /// 若尚未订阅，自动建立订阅。持锁轮询期间独占该 topic 的流（PubSub 单消费者模式）。
    async fn consume_pubsub(&self, topic: &str) -> Result<Option<Message>, MqError> {
        // 首次消费自动建立订阅
        let need_subscribe = !self.pubsub_streams.lock().await.contains_key(topic);
        if need_subscribe {
            self.subscribe_pubsub(topic).await?;
        }
        // 持锁轮询：100ms 超时无消息返回 None
        let mut streams = self.pubsub_streams.lock().await;
        let stream = match streams.get_mut(topic) {
            Some(s) => s,
            None => {
                return Err(MqError::Subscribe(
                    "Redis pubsub stream not available".into(),
                ))
            }
        };
        match tokio::time::timeout(
            std::time::Duration::from_millis(POLL_TIMEOUT_MS),
            stream.next(),
        )
        .await
        {
            Ok(Some(msg)) => Ok(Some(Message {
                topic: msg.get_channel_name().to_string(),
                payload: msg.get_payload_bytes().to_vec(),
                key: None,
                timestamp: current_timestamp_millis(),
                headers: HashMap::new(),
                id: format!("redis-pubsub-{}", current_timestamp_millis()),
                retry_count: 0,
            })),
            _ => Ok(None),
        }
    }

    /// PubSub 模式订阅：建立独立连接、`SUBSCRIBE` channel、`into_on_message` 得到 owned stream
    async fn subscribe_pubsub(&self, topic: &str) -> Result<(), MqError> {
        let url = self
            .config
            .url
            .clone()
            .unwrap_or_else(|| "redis://127.0.0.1:6379/0".to_string());
        let client = redis::Client::open(url.as_str())
            .map_err(|e| MqError::Connection(format!("Redis client open failed: {e}")))?;
        let mut pubsub = client
            .get_async_pubsub()
            .await
            .map_err(|e| MqError::Connection(format!("Redis pubsub connect failed: {e}")))?;
        pubsub
            .subscribe(topic)
            .await
            .map_err(|e| MqError::Subscribe(format!("Redis SUBSCRIBE failed: {e}")))?;
        let stream = pubsub.into_on_message();
        self.pubsub_streams
            .lock()
            .await
            .insert(topic.to_string(), Box::pin(stream));
        Ok(())
    }
    // ---------------- Stream 模式 ----------------

    /// Stream 模式发布：`XADD` 写入一条 entry（字段 `payload` 存原始字节），返回 entry id
    async fn publish_stream(&self, topic: &str, message: &[u8]) -> Result<(), MqError> {
        let mut conn = self.conn.clone();
        let field_value = ("payload", message.to_vec());
        let _: String = conn
            .xadd(topic, "*", &[field_value])
            .await
            .map_err(|e| MqError::Publish(format!("Redis XADD failed: {e}")))?;
        Ok(())
    }

    /// Stream 模式消费：`XREADGROUP` 从消费者组拉取一条新消息（`">"`）
    ///
    /// 将 `"{topic}::{entry_id}"` 记入 in_flight 表，ack 时据此调用 `XACK`。
    async fn consume_stream(&self, topic: &str) -> Result<Option<Message>, MqError> {
        // 首次消费前确保消费者组存在（已存在则忽略 BUSYGROUP）
        self.ensure_group(topic).await?;

        let mut conn = self.conn.clone();
        let options = redis::streams::StreamReadOptions::default()
            .group(self.group(), self.consumer())
            .count(1)
            .block(POLL_TIMEOUT_MS as usize);
        let keys = [topic];
        let ids = [">"];
        let reply: redis::streams::StreamReadReply = conn
            .xread_options(&keys[..], &ids[..], &options)
            .await
            .map_err(|e| MqError::Connection(format!("Redis XREADGROUP failed: {e}")))?;

        for stream_key in reply.keys {
            for entry in stream_key.ids {
                let entry_id = entry.id.clone();
                let payload: Vec<u8> = entry.get("payload").unwrap_or_default();
                let message_id = format!("{topic}::{entry_id}");
                // 记录元信息供 ack 使用
                self.in_flight
                    .lock()
                    .await
                    .insert(message_id.clone(), (topic.to_string(), self.group().to_string()));
                return Ok(Some(Message {
                    topic: topic.to_string(),
                    payload,
                    key: Some(entry_id),
                    timestamp: current_timestamp_millis(),
                    headers: HashMap::new(),
                    id: message_id,
                    retry_count: 0,
                }));
            }
        }
        Ok(None)
    }

    /// 确保消费者组存在（已存在则忽略 BUSYGROUP 错误）
    async fn ensure_group(&self, topic: &str) -> Result<(), MqError> {
        let mut conn = self.conn.clone();
        let result: Result<(), RedisError> = conn
            .xgroup_create(topic, self.group(), STREAM_START_FROM_BEGINNING)
            .await;
        if let Err(e) = result {
            let msg = e.to_string();
            // BUSYGROUP 表示组已存在，可安全忽略
            if !msg.contains("BUSYGROUP") {
                return Err(MqError::Connection(format!(
                    "Redis XGROUP CREATE failed: {msg}"
                )));
            }
        }
        Ok(())
    }
}
#[async_trait]
impl MessageQueue for RedisQueueProvider {
    async fn publish(&self, topic: &str, message: &[u8]) -> Result<(), MqError> {
        match self.config.mode {
            RedisMode::List => self.publish_list(topic, message).await,
            RedisMode::PubSub => self.publish_pubsub(topic, message).await,
            RedisMode::Stream => self.publish_stream(topic, message).await,
        }
    }

    async fn consume(&self, topic: &str) -> Result<Option<Message>, MqError> {
        match self.config.mode {
            RedisMode::List => self.consume_list(topic).await,
            RedisMode::PubSub => self.consume_pubsub(topic).await,
            RedisMode::Stream => self.consume_stream(topic).await,
        }
    }

    /// 确认消息已处理
    ///
    /// - **List / PubSub**：无 ACK 概念，返回 `Ok(())`（消费即确认，幂等）
    /// - **Stream**：根据 in_flight 表中记录的 (stream, group)，调用 `XACK` 确认该 entry
    async fn ack(&self, message_id: &str) -> Result<(), MqError> {
        let entry = self.in_flight.lock().await.remove(message_id);
        match entry {
            None => Ok(()), // List/PubSub 模式或未知 id：幂等视为已确认
            Some((stream, group)) => {
                // message_id 形如 "{topic}::{entry_id}"，提取 entry_id 部分
                let entry_id = message_id
                    .rsplit_once("::")
                    .map(|(_, id)| id.to_string())
                    .unwrap_or_else(|| message_id.to_string());
                let mut conn = self.conn.clone();
                let _: i64 = conn
                    .xack(stream, group, &[entry_id])
                    .await
                    .map_err(|e| MqError::Publish(format!("Redis XACK failed: {e}")))?;
                Ok(())
            }
        }
    }

    /// 订阅 topic
    ///
    /// - **PubSub**：建立 `SUBSCRIBE` 订阅
    /// - **Stream**：确保消费者组存在（首次消费时也会自动确保）
    /// - **List**：无需订阅，返回 `Ok(())`
    async fn subscribe(&self, topic: &str) -> Result<(), MqError> {
        match self.config.mode {
            RedisMode::PubSub => self.subscribe_pubsub(topic).await,
            RedisMode::Stream => self.ensure_group(topic).await,
            RedisMode::List => Ok(()),
        }
    }
}

/// 当前时间戳（毫秒）
fn current_timestamp_millis() -> i64 {
    use std::time::{SystemTime, UNIX_EPOCH};
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_millis() as i64
}
#[cfg(test)]
mod tests {
    use super::*;

    /// 编译验证：确保 Redis 配置/模式类型在 feature 启用时编译通过
    #[test]
    fn test_redis_config_types_compile() {
        let cfg = RedisConfig {
            mode: RedisMode::Stream,
            ..RedisConfig::default()
        };
        assert_eq!(cfg.mode, RedisMode::Stream);
        assert_eq!(cfg.pool_size, 8);
    }

    #[test]
    fn test_redis_config_default_is_list_mode() {
        let cfg = RedisConfig::default();
        assert_eq!(cfg.mode, RedisMode::List);
        assert!(cfg.consumer_group.is_none());
    }

    // ----------------------------------------------------------------------
    // 真实 Redis 集成测试（需启动 Redis 服务器，默认 #[ignore]）
    // 启动方式：docker run -d -p 6379:6379 redis:7
    // ----------------------------------------------------------------------

    #[tokio::test]
    #[ignore = "需真实 Redis 服务器"]
    async fn test_redis_list_publish_and_consume() {
        let queue = RedisQueueProvider::connect(RedisConfig::default())
            .await
            .unwrap();
        queue.publish("sz-orm-list", b"hello-list").await.unwrap();
        let msg = queue
            .consume("sz-orm-list")
            .await
            .unwrap()
            .expect("应有消息");
        assert_eq!(msg.payload, b"hello-list");
        // List 模式 ack 为 no-op
        queue.ack(&msg.id).await.unwrap();
    }

    #[tokio::test]
    #[ignore = "需真实 Redis 服务器"]
    async fn test_redis_pubsub_publish_and_consume() {
        let queue = RedisQueueProvider::connect(RedisConfig {
            mode: RedisMode::PubSub,
            ..RedisConfig::default()
        })
        .await
        .unwrap();
        // 先订阅并等待订阅生效
        queue.subscribe("sz-orm-pubsub").await.unwrap();
        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
        queue
            .publish("sz-orm-pubsub", b"hello-pubsub")
            .await
            .unwrap();
        let msg = queue
            .consume("sz-orm-pubsub")
            .await
            .unwrap()
            .expect("应有消息");
        assert_eq!(msg.payload, b"hello-pubsub");
    }

    #[tokio::test]
    #[ignore = "需真实 Redis 服务器"]
    async fn test_redis_stream_publish_consume_ack() {
        let queue = RedisQueueProvider::connect(RedisConfig {
            mode: RedisMode::Stream,
            ..RedisConfig::default()
        })
        .await
        .unwrap();
        // 清理可能残留的 stream（保证测试幂等）
        {
            let mut conn = queue.conn.clone();
            let _: i64 = conn.del("sz-orm-stream").await.unwrap_or(0);
        }
        queue
            .publish("sz-orm-stream", b"hello-stream")
            .await
            .unwrap();
        let msg = queue
            .consume("sz-orm-stream")
            .await
            .unwrap()
            .expect("应有消息");
        assert_eq!(msg.payload, b"hello-stream");
        // Stream 模式 ack 应真正调用 XACK
        queue.ack(&msg.id).await.unwrap();
    }
}