PicoBot/src/gateway/session_pool.rs

use std::collections::HashMap;
use std::sync::Arc;
use std::time::Instant;

use tokio::sync::{Mutex, mpsc};

use crate::agent::AgentError;
use crate::protocol::WsOutbound;

use super::session::Session;
use super::session_factory::SessionFactory;

/// 判断 chat_id 是否是定时任务专用（以 "scheduler/" 开头）
pub(crate) fn is_scheduler_chat_id(chat_id: &str) -> bool {
    chat_id.starts_with("scheduler/")
}

#[derive(Clone)]
pub(crate) struct SessionPool {
    inner: Arc<Mutex<SessionPoolInner>>,
    session_factory: SessionFactory,
    session_ttl_hours: Option<u64>,
}

struct SessionPoolInner {
    /// 主 Session：用于用户消息
    sessions: HashMap<String, Arc<Mutex<Session>>>,
    /// 定时任务专用 Session：独立的实例，避免与用户消息竞争锁
    scheduler_sessions: HashMap<String, Arc<Mutex<Session>>>,
    session_timestamps: HashMap<String, Instant>,
}

impl SessionPool {
    pub(crate) fn new(session_factory: SessionFactory, session_ttl_hours: Option<u64>) -> Self {
        Self {
            inner: Arc::new(Mutex::new(SessionPoolInner {
                sessions: HashMap::new(),
                scheduler_sessions: HashMap::new(),
                session_timestamps: HashMap::new(),
            })),
            session_factory,
            session_ttl_hours,
        }
    }

    /// 确保主 Session 存在（用于用户消息）
    pub(crate) async fn ensure_session(&self, channel_name: &str) -> Result<(), AgentError> {
        self.ensure_session_internal(channel_name, false).await
    }

    /// 确保定时任务专用 Session 存在
    pub(crate) async fn ensure_scheduler_session(
        &self,
        channel_name: &str,
    ) -> Result<(), AgentError> {
        self.ensure_session_internal(channel_name, true).await
    }

    /// 内部方法：创建 Session（根据 is_scheduler 选择存储位置）
    ///
    /// 使用 double-checked locking：先短暂持锁检查存在性，释放锁后执行耗时的
    /// session 创建（含配置加载、agent 工厂构造），再次持锁插入并处理竞态。
    /// 避免跨 `session_factory.create().await` 持有全局锁导致所有 channel 的
    /// session 访问串行化。
    async fn ensure_session_internal(
        &self,
        channel_name: &str,
        is_scheduler: bool,
    ) -> Result<(), AgentError> {
        // Fast path: 已存在直接返回（短暂持锁）
        {
            let inner = self.inner.lock().await;
            let sessions = if is_scheduler {
                &inner.scheduler_sessions
            } else {
                &inner.sessions
            };
            if sessions.contains_key(channel_name) {
                return Ok(());
            }
        }

        // Slow path: 锁外创建 session（耗时操作：加载配置、构造 agent 工厂）
        let (user_tx, _rx) = mpsc::channel::<WsOutbound>(100);
        let session = self
            .session_factory
            .create(channel_name.to_string(), user_tx)
            .await?;

        // 再次持锁插入，处理竞态（另一个并发任务可能已插入）
        let mut inner = self.inner.lock().await;
        let sessions = if is_scheduler {
            &mut inner.scheduler_sessions
        } else {
            &mut inner.sessions
        };
        if sessions.contains_key(channel_name) {
            // 竞态：另一任务先插入，丢弃我们创建的 session
            // （drop session 释放资源，user_tx 也 drop，无泄漏）
            tracing::debug!(
                channel = %channel_name,
                "Session created concurrently by another task, discarding duplicate"
            );
            return Ok(());
        }
        sessions.insert(channel_name.to_string(), Arc::new(Mutex::new(session)));
        inner
            .session_timestamps
            .insert(channel_name.to_string(), Instant::now());
        Ok(())
    }

    /// 获取主 Session（用于用户消息）
    pub(crate) async fn get(&self, channel_name: &str) -> Option<Arc<Mutex<Session>>> {
        self.inner.lock().await.sessions.get(channel_name).cloned()
    }

    /// 获取定时任务专用 Session
    pub(crate) async fn get_scheduler_session(
        &self,
        channel_name: &str,
    ) -> Option<Arc<Mutex<Session>>> {
        self.inner
            .lock()
            .await
            .scheduler_sessions
            .get(channel_name)
            .cloned()
    }

    /// 根据 chat_id 自动选择 Session
    /// - scheduler/ 开头：返回定时任务专用 Session
    /// - 其他：返回主 Session
    pub(crate) async fn get_for_chat_id(
        &self,
        channel_name: &str,
        chat_id: &str,
    ) -> Option<Arc<Mutex<Session>>> {
        if is_scheduler_chat_id(chat_id) {
            self.get_scheduler_session(channel_name).await
        } else {
            self.get(channel_name).await
        }
    }

    /// 确保 Session 存在（根据 chat_id 自动选择）
    pub(crate) async fn ensure_session_for_chat_id(
        &self,
        channel_name: &str,
        chat_id: &str,
    ) -> Result<(), AgentError> {
        if is_scheduler_chat_id(chat_id) {
            self.ensure_scheduler_session(channel_name).await
        } else {
            self.ensure_session(channel_name).await
        }
    }

    pub(crate) async fn touch(&self, channel_name: &str) {
        self.inner
            .lock()
            .await
            .session_timestamps
            .insert(channel_name.to_string(), Instant::now());
    }

    pub(crate) async fn cleanup_expired_sessions(&self) -> usize {
        let ttl_hours = match self.session_ttl_hours {
            Some(hours) if hours > 0 => hours,
            _ => return 0,
        };

        let ttl_duration = std::time::Duration::from_secs(ttl_hours * 3600);
        let mut inner = self.inner.lock().await;
        let now = Instant::now();

        let expired_channels: Vec<String> = inner
            .session_timestamps
            .iter()
            .filter_map(|(channel_name, last_active)| {
                let elapsed = now.duration_since(*last_active);
                if elapsed >= ttl_duration {
                    tracing::info!(
                        channel = %channel_name,
                        elapsed_hours = elapsed.as_secs() / 3600,
                        ttl_hours = ttl_hours,
                        "Session expired, removing from memory pool"
                    );
                    Some(channel_name.clone())
                } else {
                    None
                }
            })
            .collect();

        for channel_name in &expired_channels {
            // 清理主 Session
            inner.sessions.remove(channel_name);
            // 清理定时任务专用 Session
            inner.scheduler_sessions.remove(channel_name);
            inner.session_timestamps.remove(channel_name);
        }

        expired_channels.len()
    }
}