hikari-desktop/src-tauri/src/bridge_manager.rs

use parking_lot::Mutex;
use std::collections::HashMap;
use std::sync::Arc;
use tauri::AppHandle;

use crate::commands::record_session;
use crate::config::ClaudeStartOptions;
use crate::stats::UsageStats;
use crate::wsl_bridge::WslBridge;

pub struct BridgeManager {
    bridges: HashMap<String, WslBridge>,
    app_handle: Option<AppHandle>,
}

impl BridgeManager {
    pub fn new() -> Self {
        BridgeManager {
            bridges: HashMap::new(),
            app_handle: None,
        }
    }

    pub fn set_app_handle(&mut self, app: AppHandle) {
        self.app_handle = Some(app);
    }

    pub fn start_claude(
        &mut self,
        conversation_id: &str,
        options: ClaudeStartOptions,
    ) -> Result<(), String> {
        // Check if a bridge already exists and is running for this conversation
        if self
            .bridges
            .get(conversation_id)
            .map(|b| b.is_running())
            .unwrap_or(false)
        {
            return Err("Claude is already running for this conversation".to_string());
        }

        let app = self
            .app_handle
            .as_ref()
            .ok_or_else(|| "App handle not set".to_string())?
            .clone();

        // Reuse existing bridge if it exists (preserves stats across reconnects)
        // Only create a new bridge if one doesn't exist for this conversation
        let bridge = self
            .bridges
            .entry(conversation_id.to_string())
            .or_insert_with(|| WslBridge::new_with_conversation_id(conversation_id.to_string()));

        // Start the Claude process
        bridge.start(app.clone(), options)?;

        // Record session start for cost tracking
        tauri::async_runtime::spawn(async move {
            record_session(&app).await;
        });

        Ok(())
    }

    pub fn stop_claude(&mut self, conversation_id: &str) -> Result<(), String> {
        if let Some(bridge) = self.bridges.get_mut(conversation_id) {
            let app = self
                .app_handle
                .as_ref()
                .ok_or_else(|| "App handle not set".to_string())?;
            bridge.stop(app);
            Ok(())
        } else {
            Err("No Claude instance found for this conversation".to_string())
        }
    }

    pub fn interrupt_claude(&mut self, conversation_id: &str) -> Result<(), String> {
        if let Some(bridge) = self.bridges.get_mut(conversation_id) {
            let app = self
                .app_handle
                .as_ref()
                .ok_or_else(|| "App handle not set".to_string())?;
            bridge.interrupt(app)
        } else {
            Err("No Claude instance found for this conversation".to_string())
        }
    }

    pub fn send_prompt(&mut self, conversation_id: &str, message: String) -> Result<(), String> {
        if let Some(bridge) = self.bridges.get_mut(conversation_id) {
            bridge.send_message(&message)
        } else {
            Err("No Claude instance found for this conversation".to_string())
        }
    }

    pub fn send_tool_result(
        &mut self,
        conversation_id: &str,
        tool_use_id: &str,
        result: serde_json::Value,
    ) -> Result<(), String> {
        if let Some(bridge) = self.bridges.get_mut(conversation_id) {
            bridge.send_tool_result(tool_use_id, result)
        } else {
            Err("No Claude instance found for this conversation".to_string())
        }
    }

    pub fn is_claude_running(&self, conversation_id: &str) -> bool {
        self.bridges
            .get(conversation_id)
            .map(|b| b.is_running())
            .unwrap_or(false)
    }

    pub fn get_working_directory(&self, conversation_id: &str) -> Result<String, String> {
        self.bridges
            .get(conversation_id)
            .map(|b| b.get_working_directory().to_string())
            .ok_or_else(|| "No Claude instance found for this conversation".to_string())
    }

    pub fn get_usage_stats(&self, conversation_id: &str) -> Result<UsageStats, String> {
        self.bridges
            .get(conversation_id)
            .map(|b| b.get_stats())
            .ok_or_else(|| "No Claude instance found for this conversation".to_string())
    }

    #[allow(dead_code)]
    pub fn cleanup_stopped_bridges(&mut self) {
        // Remove bridges that are no longer running
        self.bridges.retain(|_, bridge| bridge.is_running());
    }

    #[allow(dead_code)]
    pub fn stop_all(&mut self) {
        if let Some(app) = &self.app_handle {
            for (_, bridge) in self.bridges.iter_mut() {
                bridge.stop(app);
            }
        }
        self.bridges.clear();
    }

    #[allow(dead_code)]
    pub fn get_active_conversations(&self) -> Vec<String> {
        self.bridges
            .keys()
            .filter(|id| {
                self.bridges
                    .get(*id)
                    .map(|b| b.is_running())
                    .unwrap_or(false)
            })
            .cloned()
            .collect()
    }
}

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

pub type SharedBridgeManager = Arc<Mutex<BridgeManager>>;

pub fn create_shared_bridge_manager() -> SharedBridgeManager {
    Arc::new(Mutex::new(BridgeManager::new()))
}

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

    #[test]
    fn test_bridge_manager_new() {
        let manager = BridgeManager::new();
        assert!(manager.app_handle.is_none());
        assert!(manager.bridges.is_empty());
    }

    #[test]
    fn test_bridge_manager_default() {
        let manager = BridgeManager::default();
        assert!(manager.app_handle.is_none());
        assert!(manager.bridges.is_empty());
    }

    #[test]
    fn test_is_claude_running_no_bridge() {
        let manager = BridgeManager::new();
        assert!(!manager.is_claude_running("nonexistent"));
    }

    #[test]
    fn test_get_working_directory_no_bridge() {
        let manager = BridgeManager::new();
        let result = manager.get_working_directory("nonexistent");
        assert!(result.is_err());
        assert_eq!(
            result.unwrap_err(),
            "No Claude instance found for this conversation"
        );
    }

    #[test]
    fn test_get_usage_stats_no_bridge() {
        let manager = BridgeManager::new();
        let result = manager.get_usage_stats("nonexistent");
        assert!(result.is_err());
        assert_eq!(
            result.unwrap_err(),
            "No Claude instance found for this conversation"
        );
    }

    #[test]
    fn test_stop_claude_no_bridge() {
        let mut manager = BridgeManager::new();
        let result = manager.stop_claude("nonexistent");
        assert!(result.is_err());
        assert_eq!(
            result.unwrap_err(),
            "No Claude instance found for this conversation"
        );
    }

    #[test]
    fn test_interrupt_claude_no_bridge() {
        let mut manager = BridgeManager::new();
        let result = manager.interrupt_claude("nonexistent");
        assert!(result.is_err());
        assert_eq!(
            result.unwrap_err(),
            "No Claude instance found for this conversation"
        );
    }

    #[test]
    fn test_send_prompt_no_bridge() {
        let mut manager = BridgeManager::new();
        let result = manager.send_prompt("nonexistent", "Hello".to_string());
        assert!(result.is_err());
        assert_eq!(
            result.unwrap_err(),
            "No Claude instance found for this conversation"
        );
    }

    #[test]
    fn test_send_tool_result_no_bridge() {
        let mut manager = BridgeManager::new();
        let result = manager.send_tool_result(
            "nonexistent",
            "tool_id",
            serde_json::json!({"result": "success"}),
        );
        assert!(result.is_err());
        assert_eq!(
            result.unwrap_err(),
            "No Claude instance found for this conversation"
        );
    }

    #[test]
    fn test_create_shared_bridge_manager() {
        let shared = create_shared_bridge_manager();
        let manager = shared.lock();
        assert!(manager.bridges.is_empty());
        assert!(manager.app_handle.is_none());
    }

    #[test]
    fn test_cleanup_stopped_bridges_empty() {
        let mut manager = BridgeManager::new();
        manager.cleanup_stopped_bridges();
        assert!(manager.bridges.is_empty());
    }

    #[test]
    fn test_get_active_conversations_empty() {
        let manager = BridgeManager::new();
        let active = manager.get_active_conversations();
        assert!(active.is_empty());
    }

    #[test]
    fn test_stop_all_without_app_handle() {
        let mut manager = BridgeManager::new();
        manager.stop_all(); // Should not panic
        assert!(manager.bridges.is_empty());
    }
}