QuiCommit/src/llm/thinking.rs

use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};

/// 统一的思考状态管理器，用于管理模型思考状态的显示与隐藏
pub struct ThinkingStateManager {
    is_thinking: AtomicBool,
    on_start: Option<Box<dyn Fn() + Send + Sync>>,
    on_end: Option<Box<dyn Fn() + Send + Sync>>,
}

impl ThinkingStateManager {
    pub fn new() -> Self {
        Self {
            is_thinking: AtomicBool::new(false),
            on_start: None,
            on_end: None,
        }
    }

    /// 设置思考开始回调
    pub fn on_thinking_start<F: Fn() + Send + Sync + 'static>(mut self, callback: F) -> Self {
        self.on_start = Some(Box::new(callback));
        self
    }

    /// 设置思考结束回调
    pub fn on_thinking_end<F: Fn() + Send + Sync + 'static>(mut self, callback: F) -> Self {
        self.on_end = Some(Box::new(callback));
        self
    }

    /// 开始思考状态
    pub fn start_thinking(&self) {
        if !self.is_thinking.load(Ordering::SeqCst) {
            self.is_thinking.store(true, Ordering::SeqCst);
            if let Some(ref cb) = self.on_start {
                cb();
            }
        }
    }

    /// 结束思考状态
    pub fn end_thinking(&self) {
        if self.is_thinking.load(Ordering::SeqCst) {
            self.is_thinking.store(false, Ordering::SeqCst);
            if let Some(ref cb) = self.on_end {
                cb();
            }
        }
    }

    /// 当前是否处于思考状态
    pub fn is_thinking(&self) -> bool {
        self.is_thinking.load(Ordering::SeqCst)
    }
}

impl Default for ThinkingStateManager {
    fn default() -> Self {
        Self::new()
    }
}

/// 线程安全的思考状态管理器引用
pub type SharedThinkingState = Arc<ThinkingStateManager>;

/// 创建 LLM 流式使用的共享思考状态。
/// 进度显示由 AI 生成 spinner 负责（issue 21），此处不再附加控制台输出。
pub fn create_console_thinking_state() -> SharedThinkingState {
    Arc::new(ThinkingStateManager::new())
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::Mutex;

    #[test]
    fn test_thinking_state_transitions() {
        let manager = ThinkingStateManager::new();
        assert!(!manager.is_thinking());

        manager.start_thinking();
        assert!(manager.is_thinking());

        manager.end_thinking();
        assert!(!manager.is_thinking());
    }

    #[test]
    fn test_thinking_idempotent_start() {
        let manager = ThinkingStateManager::new();
        manager.start_thinking();
        manager.start_thinking(); // 重复调用不应触发回调两次
        assert!(manager.is_thinking());
    }

    #[test]
    fn test_thinking_idempotent_end() {
        let manager = ThinkingStateManager::new();
        manager.end_thinking(); // 未开始时结束不应触发问题
        assert!(!manager.is_thinking());
    }

    #[test]
    fn test_thinking_callbacks() {
        let events: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
        let events_clone = events.clone();

        let manager = ThinkingStateManager::new().on_thinking_start(move || {
            events_clone.lock().unwrap().push("start".to_string());
        });

        let events_clone2 = events.clone();
        let manager = manager.on_thinking_end(move || {
            events_clone2.lock().unwrap().push("end".to_string());
        });

        manager.start_thinking();
        manager.end_thinking();

        let recorded = events.lock().unwrap();
        assert_eq!(recorded.len(), 2);
        assert_eq!(recorded[0], "start");
        assert_eq!(recorded[1], "end");
    }

    #[test]
    fn test_create_console_thinking_state() {
        let state = create_console_thinking_state();
        assert!(!state.is_thinking());
        state.start_thinking();
        assert!(state.is_thinking());
        state.end_thinking();
        assert!(!state.is_thinking());
    }

    #[test]
    fn test_default() {
        let manager = ThinkingStateManager::default();
        assert!(!manager.is_thinking());
    }
}