gitdataai/libs/db/cache.rs

use config::AppConfig;
use deadpool_redis::cluster::{Connection, Manager, Pool};
use uuid::Uuid;

const CHAT_STREAM_KEY_PREFIX: &str = "chat:stream:";
const CHAT_STREAM_TTL_SECS: u64 = 600; // 10 minutes

#[derive(Clone)]
pub struct AppCache {
    pool: Pool,
    redis_url: String,
}

impl AppCache {
    pub async fn init(cfg: &AppConfig) -> anyhow::Result<Self> {
        let urls = cfg.redis_urls()?;
        let pool_size = cfg.redis_pool_size()?;
        let conn = Manager::new(urls.clone(), true)?;
        let pool = deadpool_redis::cluster::Pool::builder(conn)
            .max_size(pool_size as usize)
            .build()?;
        let redis_url = urls
            .first()
            .cloned()
            .unwrap_or_else(|| "redis://127.0.0.1:6379".to_string());

        Ok(Self { pool, redis_url })
    }
    pub async fn conn(&self) -> anyhow::Result<Connection> {
        Ok(self.pool.get().await?)
    }

    pub fn redis_pool(&self) -> &Pool {
        &self.pool
    }

    pub fn redis_url(&self) -> &str {
        &self.redis_url
    }

    /// Set chat stream active state (conversation_id → {message_id, started_at}).
    /// TTL 10 minutes — prevents stale entries.
    pub async fn set_chat_stream_active(&self, conversation_id: Uuid, message_id: Uuid) -> bool {
        if let Ok(mut conn) = self.conn().await {
            let key = format!("{}{}", CHAT_STREAM_KEY_PREFIX, conversation_id);
            let value = serde_json::json!({
                "message_id": message_id.to_string(),
                "started_at": chrono::Utc::now().timestamp(),
            })
            .to_string();
            let _: Result<(), _> = redis::cmd("SETEX")
                .arg(&key)
                .arg(CHAT_STREAM_TTL_SECS as i64)
                .arg(&value)
                .query_async(&mut conn)
                .await;
            return true;
        }
        false
    }

    /// Get active chat stream state for a conversation. Returns (message_id, started_at_ts).
    pub async fn get_chat_stream_active(&self, conversation_id: Uuid) -> Option<(Uuid, i64)> {
        if let Ok(mut conn) = self.conn().await {
            let key = format!("{}{}", CHAT_STREAM_KEY_PREFIX, conversation_id);
            if let Ok(value) = redis::cmd("GET")
                .arg(&key)
                .query_async::<String>(&mut conn)
                .await
            {
                if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(&value) {
                    let msg_id = parsed
                        .get("message_id")
                        .and_then(|v| v.as_str())
                        .and_then(|s| Uuid::parse_str(s).ok());
                    let started_at = parsed.get("started_at").and_then(|v| v.as_i64());
                    if let (Some(mid), Some(ts)) = (msg_id, started_at) {
                        return Some((mid, ts));
                    }
                }
            }
        }
        None
    }

    /// Clear chat stream active state (called when streaming finishes).
    pub async fn clear_chat_stream_active(&self, conversation_id: Uuid) {
        if let Ok(mut conn) = self.conn().await {
            let key = format!("{}{}", CHAT_STREAM_KEY_PREFIX, conversation_id);
            let _: Result<(), _> = redis::cmd("DEL").arg(&key).query_async(&mut conn).await;
        }
    }

    // ─── Chat Stream Cancellation ─────────────────────────────────────────

    /// Mark a chat stream as cancelled by the user.
    /// TTL 60 seconds — prevents stale entries.
    pub async fn set_chat_stream_cancelled(&self, conversation_id: Uuid) -> bool {
        if let Ok(mut conn) = self.conn().await {
            let key = format!("{}cancel:{}", CHAT_STREAM_KEY_PREFIX, conversation_id);
            let _: Result<(), _> = redis::cmd("SETEX")
                .arg(&key)
                .arg(60_i64)
                .arg("1")
                .query_async(&mut conn)
                .await;
            return true;
        }
        false
    }

    /// Check if a chat stream has been cancelled.
    pub async fn is_chat_stream_cancelled(&self, conversation_id: Uuid) -> bool {
        if let Ok(mut conn) = self.conn().await {
            let key = format!("{}cancel:{}", CHAT_STREAM_KEY_PREFIX, conversation_id);
            if let Ok(value) = redis::cmd("GET")
                .arg(&key)
                .query_async::<Option<String>>(&mut conn)
                .await
            {
                return value.is_some();
            }
        }
        false
    }

    /// Clear the cancel flag for a chat stream.
    pub async fn clear_chat_stream_cancelled(&self, conversation_id: Uuid) {
        if let Ok(mut conn) = self.conn().await {
            let key = format!("{}cancel:{}", CHAT_STREAM_KEY_PREFIX, conversation_id);
            let _: Result<(), _> = redis::cmd("DEL").arg(&key).query_async(&mut conn).await;
        }
    }

    pub async fn set_sub_agent_cancelled(&self, conversation_id: Uuid, children_id: &str) -> bool {
        if let Ok(mut conn) = self.conn().await {
            let key = format!(
                "{}subagent:cancel:{}:{}",
                CHAT_STREAM_KEY_PREFIX, conversation_id, children_id
            );
            let _: Result<(), _> = redis::cmd("SETEX")
                .arg(&key)
                .arg(300_i64)
                .arg("1")
                .query_async(&mut conn)
                .await;
            return true;
        }
        false
    }

    pub async fn is_sub_agent_cancelled(&self, conversation_id: Uuid, children_id: &str) -> bool {
        if let Ok(mut conn) = self.conn().await {
            let key = format!(
                "{}subagent:cancel:{}:{}",
                CHAT_STREAM_KEY_PREFIX, conversation_id, children_id
            );
            if let Ok(value) = redis::cmd("GET")
                .arg(&key)
                .query_async::<Option<String>>(&mut conn)
                .await
            {
                return value.is_some();
            }
        }
        false
    }

    pub async fn clear_sub_agent_cancelled(&self, conversation_id: Uuid, children_id: &str) {
        if let Ok(mut conn) = self.conn().await {
            let key = format!(
                "{}subagent:cancel:{}:{}",
                CHAT_STREAM_KEY_PREFIX, conversation_id, children_id
            );
            let _: Result<(), _> = redis::cmd("DEL").arg(&key).query_async(&mut conn).await;
        }
    }
}