mcp-stdio-proxy 0.2.2

use axum::Router;
use dashmap::DashMap;
use log::{debug, error, info};
use moka::future::Cache;
use once_cell::sync::Lazy;
use std::sync::Arc;
use tokio::sync::{Mutex, OwnedMutexGuard};
use tokio::time::{Duration, Instant};
use tokio_util::sync::CancellationToken;
use tracing::warn;

use anyhow::Result;

use crate::proxy::McpHandler;

use super::{CheckMcpStatusResponseStatus, McpConfig, McpProtocol, McpRouterPath, McpType};

// 全局单例路由表
pub static GLOBAL_ROUTES: Lazy<Arc<DashMap<String, Router>>> =
    Lazy::new(|| Arc::new(DashMap::new()));

// 全局单例 ProxyHandlerManager
pub static GLOBAL_PROXY_MANAGER: Lazy<ProxyHandlerManager> =
    Lazy::new(ProxyHandlerManager::default);

/// 动态路由服务
#[derive(Clone)]
pub struct DynamicRouterService(pub McpProtocol);

impl DynamicRouterService {
    // 注册动态 handler
    pub fn register_route(path: &str, handler: Router) {
        debug!("=== Register Route ===");
        debug!("Registration path: {}", path);
        GLOBAL_ROUTES.insert(path.to_string(), handler);
        debug!("=== Route registration completed ===");
    }

    // 删除动态 handler
    pub fn delete_route(path: &str) {
        debug!("=== Delete route ===");
        debug!("Delete path: {}", path);
        GLOBAL_ROUTES.remove(path);
        debug!("=== Route deletion completed ===");
    }

    // 获取动态 handler
    pub fn get_route(path: &str) -> Option<Router> {
        let result = GLOBAL_ROUTES.get(path).map(|entry| entry.value().clone());
        if result.is_some() {
            debug!("get_route('{}') = Some(Router)", path);
        } else {
            debug!("get_route('{}') = None", path);
        }
        result
    }

    // 获取所有已注册的路由（debug用）
    pub fn get_all_routes() -> Vec<String> {
        GLOBAL_ROUTES
            .iter()
            .map(|entry| entry.key().clone())
            .collect()
    }
}

impl std::fmt::Debug for DynamicRouterService {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let routes = GLOBAL_ROUTES
            .iter()
            .map(|entry| entry.key().clone())
            .collect::<Vec<_>>();
        write!(f, "DynamicRouterService {{ routes: {routes:?} }}")
    }
}

// =============================================================================
// RAII 进程管理设计
// =============================================================================
//
// 设计目标：当 mcp_id 从 map 中移除时，自动释放对应的 MCP 进程资源
//
// 核心结构：
// - McpProcessGuard: 进程生命周期守护器，实现 Drop trait 自动取消 CancellationToken
// - McpService: 封装 McpHandler + McpProcessGuard + 服务状态，作为 map 的 value
// - ProxyHandlerManager: 使用单一 DashMap<String, McpService> 管理所有服务
//
// 资源释放流程：
// 1. 从 map 中 remove mcp_id
// 2. McpService 被 drop
// 3. McpProcessGuard::drop() 被调用
// 4. CancellationToken 被 cancel
// 5. 监听该 token 的 SseServer/子进程收到信号，自动退出
// =============================================================================

/// MCP 进程生命周期守护器
///
/// 实现 RAII 模式：当此结构体被 drop 时，自动取消 CancellationToken，
/// 触发关联的 SseServer 和子进程退出。
///
/// # 使用场景
///
/// 1. 从 `ProxyHandlerManager` 移除 mcp_id 时，自动清理进程
/// 2. 服务重启时，旧服务自动被清理
/// 3. 系统关闭时，所有服务自动清理
pub struct McpProcessGuard {
    mcp_id: String,
    cancellation_token: CancellationToken,
}

impl McpProcessGuard {
    pub fn new(mcp_id: String, cancellation_token: CancellationToken) -> Self {
        debug!("[RAII] Create process daemon: mcp_id={}", mcp_id);
        Self {
            mcp_id,
            cancellation_token,
        }
    }

    /// 克隆 CancellationToken（用于传递给异步任务）
    pub fn clone_token(&self) -> CancellationToken {
        self.cancellation_token.clone()
    }
}

impl Drop for McpProcessGuard {
    fn drop(&mut self) {
        info!(
            "[RAII] The process daemon was dropped and canceled CancellationToken: mcp_id={}",
            self.mcp_id
        );
        self.cancellation_token.cancel();
    }
}

// McpProcessGuard 不实现 Clone，确保唯一所有权
impl std::fmt::Debug for McpProcessGuard {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("McpProcessGuard")
            .field("mcp_id", &self.mcp_id)
            .field("is_cancelled", &self.cancellation_token.is_cancelled())
            .finish()
    }
}

/// MCP 服务封装
///
/// 将 McpHandler、McpProcessGuard 和服务状态封装在一起，
/// 作为 `ProxyHandlerManager` 中 DashMap 的 value。
///
/// # RAII 保证
///
/// 当 McpService 被 drop 时：
/// 1. McpProcessGuard 被 drop，触发 CancellationToken 取消
/// 2. 关联的子进程收到信号，自动退出
pub struct McpService {
    /// 进程守护器（RAII 核心）
    process_guard: McpProcessGuard,
    /// MCP 透明代理处理器（可选，启动中时为 None）
    handler: Option<McpHandler>,
    /// 服务状态信息
    status: McpServiceStatusInfo,
}

/// MCP 服务状态信息（不包含 CancellationToken，由 McpProcessGuard 管理）
#[derive(Debug, Clone)]
pub struct McpServiceStatusInfo {
    pub mcp_id: String,
    pub mcp_type: McpType,
    pub mcp_router_path: McpRouterPath,
    pub check_mcp_status_response_status: CheckMcpStatusResponseStatus,
    pub last_accessed: Instant,
    pub mcp_config: Option<McpConfig>,
    /// 连续健康检查失败次数（用于容错机制）
    pub consecutive_probe_failures: u32,
}

impl McpServiceStatusInfo {
    pub fn new(
        mcp_id: String,
        mcp_type: McpType,
        mcp_router_path: McpRouterPath,
        check_mcp_status_response_status: CheckMcpStatusResponseStatus,
    ) -> Self {
        Self {
            mcp_id,
            mcp_type,
            mcp_router_path,
            check_mcp_status_response_status,
            last_accessed: Instant::now(),
            mcp_config: None,
            consecutive_probe_failures: 0,
        }
    }

    pub fn update_last_accessed(&mut self) {
        self.last_accessed = Instant::now();
    }

    /// 重置健康检查失败计数
    pub fn reset_probe_failures(&mut self) {
        self.consecutive_probe_failures = 0;
    }

    /// 增加健康检查失败计数，返回增加后的值
    pub fn increment_probe_failures(&mut self) -> u32 {
        self.consecutive_probe_failures += 1;
        self.consecutive_probe_failures
    }
}

impl McpService {
    /// 创建新的 MCP 服务
    ///
    /// # 参数
    /// - `mcp_id`: 服务唯一标识
    /// - `mcp_type`: 服务类型
    /// - `mcp_router_path`: 路由路径
    /// - `cancellation_token`: 用于控制进程生命周期的取消令牌
    pub fn new(
        mcp_id: String,
        mcp_type: McpType,
        mcp_router_path: McpRouterPath,
        cancellation_token: CancellationToken,
    ) -> Self {
        let process_guard = McpProcessGuard::new(mcp_id.clone(), cancellation_token);
        let status = McpServiceStatusInfo::new(
            mcp_id,
            mcp_type,
            mcp_router_path,
            CheckMcpStatusResponseStatus::Pending,
        );
        Self {
            process_guard,
            handler: None,
            status,
        }
    }

    /// 设置 MCP Handler
    pub fn set_handler(&mut self, handler: McpHandler) {
        self.handler = Some(handler);
    }

    /// 获取 MCP Handler
    pub fn handler(&self) -> Option<&McpHandler> {
        self.handler.as_ref()
    }

    /// 获取服务状态
    pub fn status(&self) -> &McpServiceStatusInfo {
        &self.status
    }

    /// 获取可变服务状态
    pub fn status_mut(&mut self) -> &mut McpServiceStatusInfo {
        &mut self.status
    }

    /// 克隆 CancellationToken
    pub fn clone_token(&self) -> CancellationToken {
        self.process_guard.clone_token()
    }

    /// 更新服务状态
    pub fn update_status(&mut self, status: CheckMcpStatusResponseStatus) {
        self.status.check_mcp_status_response_status = status;
    }

    /// 更新最后访问时间
    pub fn update_last_accessed(&mut self) {
        self.status.update_last_accessed();
    }

    /// 设置 MCP 配置
    pub fn set_mcp_config(&mut self, config: McpConfig) {
        self.status.mcp_config = Some(config);
    }
}

impl std::fmt::Debug for McpService {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("McpService")
            .field("process_guard", &self.process_guard)
            .field("handler", &self.handler.is_some())
            .field("status", &self.status)
            .finish()
    }
}

// =============================================================================
// 兼容层：保留 McpServiceStatus 以兼容现有代码
// =============================================================================

/// MCP 服务状态（兼容层）
///
/// 保留此结构体以兼容现有代码，内部委托给 McpServiceStatusInfo
#[derive(Debug, Clone)]
pub struct McpServiceStatus {
    pub mcp_id: String,
    pub mcp_type: McpType,
    pub mcp_router_path: McpRouterPath,
    pub cancellation_token: CancellationToken,
    pub check_mcp_status_response_status: CheckMcpStatusResponseStatus,
    pub last_accessed: Instant,
    pub mcp_config: Option<McpConfig>,
    /// 连续健康检查失败次数
    pub consecutive_probe_failures: u32,
}

impl McpServiceStatus {
    pub fn new(
        mcp_id: String,
        mcp_type: McpType,
        mcp_router_path: McpRouterPath,
        cancellation_token: CancellationToken,
        check_mcp_status_response_status: CheckMcpStatusResponseStatus,
    ) -> Self {
        Self {
            mcp_id,
            mcp_type,
            mcp_router_path,
            cancellation_token,
            check_mcp_status_response_status,
            last_accessed: Instant::now(),
            mcp_config: None,
            consecutive_probe_failures: 0,
        }
    }

    pub fn with_mcp_config(mut self, mcp_config: McpConfig) -> Self {
        self.mcp_config = Some(mcp_config);
        self
    }

    pub fn update_last_accessed(&mut self) {
        self.last_accessed = Instant::now();
    }
}

// =============================================================================
// ProxyHandlerManager：使用 RAII 模式管理 MCP 服务
// =============================================================================

/// MCP 代理管理器
///
/// 使用 RAII 模式管理 MCP 服务：
/// - 从 map 中移除 mcp_id 时，自动释放对应的进程资源
/// - 不需要显式调用 cleanup 方法（但仍提供显式清理接口）
#[derive(Debug)]
pub struct ProxyHandlerManager {
    /// 使用单一 DashMap 管理所有 MCP 服务（RAII 核心）
    services: DashMap<String, McpService>,
}

impl Default for ProxyHandlerManager {
    fn default() -> Self {
        ProxyHandlerManager {
            services: DashMap::new(),
        }
    }
}

impl ProxyHandlerManager {
    /// 添加 MCP 服务（RAII 模式）
    ///
    /// 使用新的 RAII 结构创建服务，当服务被移除时会自动清理资源
    pub fn add_mcp_service(
        &self,
        mcp_id: String,
        mcp_type: McpType,
        mcp_router_path: McpRouterPath,
        cancellation_token: CancellationToken,
    ) {
        let service = McpService::new(
            mcp_id.clone(),
            mcp_type,
            mcp_router_path,
            cancellation_token,
        );

        // RAII: 如果已存在同名服务，insert 会返回旧服务，旧服务被 drop 时自动清理
        if let Some(old_service) = self.services.insert(mcp_id.clone(), service) {
            info!(
                "[RAII] Overwrite existing services, old services will be automatically cleaned up: mcp_id={}",
                mcp_id
            );
            drop(old_service);
        }
    }

    /// 添加 MCP 服务状态（兼容旧 API）
    ///
    /// 保持与现有代码的兼容性，内部转换为新的 RAII 结构
    ///
    /// 注意：`last_accessed` 会被重置为当前时间（插入视为新访问）
    pub fn add_mcp_service_status_and_proxy(
        &self,
        mcp_service_status: McpServiceStatus,
        proxy_handler: Option<McpHandler>,
    ) {
        let mcp_id = mcp_service_status.mcp_id.clone();

        // 创建 McpService 使用 RAII 模式
        // 注意：last_accessed 会在 McpServiceStatusInfo::new() 中重置为 Instant::now()
        let mut service = McpService::new(
            mcp_id.clone(),
            mcp_service_status.mcp_type,
            mcp_service_status.mcp_router_path,
            mcp_service_status.cancellation_token,
        );

        // 设置初始状态
        service.status_mut().check_mcp_status_response_status =
            mcp_service_status.check_mcp_status_response_status;

        // 设置配置（如果有）
        if let Some(config) = mcp_service_status.mcp_config {
            service.set_mcp_config(config);
        }

        // 设置 handler（如果有）
        if let Some(handler) = proxy_handler {
            service.set_handler(handler);
        }

        // RAII: 如果已存在同名服务，insert 会返回旧服务，旧服务被 drop 时自动清理
        if let Some(old_service) = self.services.insert(mcp_id.clone(), service) {
            info!(
                "[RAII] Overwrite existing services, old services will be automatically cleaned up: mcp_id={}",
                mcp_id
            );
            // old_service 在此作用域结束时 drop，触发 McpProcessGuard::drop()
            drop(old_service);
        }
    }

    /// 获取所有的 MCP 服务状态（兼容旧 API）
    ///
    /// 优化：先快速收集所有 keys，然后逐个获取详细信息
    /// 避免 iter() 长时间锁住多个分片，让其他写操作有机会执行
    pub fn get_all_mcp_service_status(&self) -> Vec<McpServiceStatus> {
        // 第一步：快速收集所有 keys（只 clone String，锁持有时间短）
        let keys: Vec<String> = self
            .services
            .iter()
            .map(|entry| entry.key().clone())
            .collect();

        // 第二步：逐个获取详细信息（每次只锁一个分片）
        keys.into_iter()
            .filter_map(|mcp_id| self.get_mcp_service_status(&mcp_id))
            .collect()
    }

    /// 获取 MCP 服务状态（兼容旧 API）
    pub fn get_mcp_service_status(&self, mcp_id: &str) -> Option<McpServiceStatus> {
        self.services.get(mcp_id).map(|entry| {
            let service = entry.value();
            let status = service.status();
            McpServiceStatus {
                mcp_id: status.mcp_id.clone(),
                mcp_type: status.mcp_type.clone(),
                mcp_router_path: status.mcp_router_path.clone(),
                cancellation_token: service.clone_token(),
                check_mcp_status_response_status: status.check_mcp_status_response_status.clone(),
                last_accessed: status.last_accessed,
                mcp_config: status.mcp_config.clone(),
                consecutive_probe_failures: status.consecutive_probe_failures,
            }
        })
    }

    /// 更新最后访问时间
    ///
    /// 使用 entry API 确保原子性操作
    pub fn update_last_accessed(&self, mcp_id: &str) {
        self.services
            .entry(mcp_id.to_string())
            .and_modify(|service| service.update_last_accessed());
    }

    /// 修改 MCP 服务状态 (Ready/Pending/Error)
    ///
    /// 使用 entry API 确保原子性操作
    pub fn update_mcp_service_status(&self, mcp_id: &str, status: CheckMcpStatusResponseStatus) {
        self.services
            .entry(mcp_id.to_string())
            .and_modify(|service| service.update_status(status));
    }

    /// 获取 MCP Handler
    pub fn get_proxy_handler(&self, mcp_id: &str) -> Option<McpHandler> {
        self.services
            .get(mcp_id)
            .and_then(|entry| entry.value().handler().cloned())
    }

    /// 获取服务的 MCP 配置（用于自动重启）
    pub fn get_mcp_config(&self, mcp_id: &str) -> Option<McpConfig> {
        self.services
            .get(mcp_id)
            .and_then(|entry| entry.value().status().mcp_config.clone())
    }

    /// 添加 MCP Handler 到已存在的服务
    ///
    /// 使用 entry API 确保原子性操作
    pub fn add_proxy_handler(&self, mcp_id: &str, proxy_handler: McpHandler) {
        match self.services.entry(mcp_id.to_string()) {
            dashmap::mapref::entry::Entry::Occupied(mut entry) => {
                entry.get_mut().set_handler(proxy_handler);
            }
            dashmap::mapref::entry::Entry::Vacant(_) => {
                warn!(
                    "[RAII] Trying to add handler to non-existent service: mcp_id={}",
                    mcp_id
                );
            }
        }
    }

    /// 检查服务是否存在
    pub fn contains_service(&self, mcp_id: &str) -> bool {
        self.services.contains_key(mcp_id)
    }

    /// 获取服务数量
    pub fn service_count(&self) -> usize {
        self.services.len()
    }

    /// 注册 MCP 配置到缓存
    pub async fn register_mcp_config(&self, mcp_id: &str, config: McpConfig) {
        GLOBAL_MCP_CONFIG_CACHE
            .insert(mcp_id.to_string(), config)
            .await;
        info!("MCP configuration registered in cache: {}", mcp_id);
    }

    /// 从缓存获取 MCP 配置
    pub async fn get_mcp_config_from_cache(&self, mcp_id: &str) -> Option<McpConfig> {
        if let Some(config) = GLOBAL_MCP_CONFIG_CACHE.get(mcp_id).await {
            debug!("Get MCP configuration from cache: {}", mcp_id);
            Some(config)
        } else {
            debug!("MCP configuration not found in cache: {}", mcp_id);
            None
        }
    }

    /// 从缓存删除 MCP 配置
    pub async fn unregister_mcp_config(&self, mcp_id: &str) {
        GLOBAL_MCP_CONFIG_CACHE.invalidate(mcp_id).await;
        info!("MCP configuration removed from cache: {}", mcp_id);
    }

    /// 清理资源 (RAII 模式简化版)
    ///
    /// 通过 RAII 模式，从 DashMap 中移除服务会自动：
    /// 1. 触发 McpProcessGuard::drop()
    /// 2. 取消 CancellationToken
    /// 3. 关联的子进程收到信号退出
    ///
    /// 此方法额外清理路由和缓存
    pub async fn cleanup_resources(&self, mcp_id: &str) -> Result<()> {
        info!("[RAII] Start cleaning up resources: mcp_id={}", mcp_id);

        // 创建路径以构建要删除的路由路径
        let mcp_sse_router_path = McpRouterPath::new(mcp_id.to_string(), McpProtocol::Sse)
            .map_err(|e| {
                anyhow::anyhow!("Failed to create SSE router path for {}: {}", mcp_id, e)
            })?;
        let base_sse_path = mcp_sse_router_path.base_path;

        let mcp_stream_router_path = McpRouterPath::new(mcp_id.to_string(), McpProtocol::Stream)
            .map_err(|e| {
                anyhow::anyhow!("Failed to create Stream router path for {}: {}", mcp_id, e)
            })?;
        let base_stream_path = mcp_stream_router_path.base_path;

        // 移除相关路由
        DynamicRouterService::delete_route(&base_sse_path);
        DynamicRouterService::delete_route(&base_stream_path);

        // RAII 核心：从 DashMap 移除会触发 McpProcessGuard::drop()
        // 这会自动取消 CancellationToken，进而触发子进程退出
        if self.services.remove(mcp_id).is_some() {
            info!(
                "[RAII] The service has been removed from the map, McpProcessGuard will automatically cancel the token: mcp_id={}",
                mcp_id
            );
        } else {
            debug!(
                "[RAII] Service does not exist, skip removal: mcp_id={}",
                mcp_id
            );
        }

        // 清理配置缓存
        self.unregister_mcp_config(mcp_id).await;

        // 清理健康状态缓存
        GLOBAL_RESTART_TRACKER.clear_health_status(mcp_id);

        info!(
            "[RAII] MCP service resource cleanup completed: mcp_id={}",
            mcp_id
        );
        Ok(())
    }

    /// 系统关闭，清理所有资源
    ///
    /// RAII 模式下，清除 DashMap 会自动释放所有资源
    pub async fn cleanup_all_resources(&self) -> Result<()> {
        info!("[RAII] Start cleaning up all MCP service resources");

        // 收集所有 mcp_id
        let mcp_ids: Vec<String> = self
            .services
            .iter()
            .map(|entry| entry.key().clone())
            .collect();

        let count = mcp_ids.len();

        // 逐个清理（包括路由和缓存）
        for mcp_id in mcp_ids {
            if let Err(e) = self.cleanup_resources(&mcp_id).await {
                error!(
                    "[RAII] Failed to clean up resources: mcp_id={}, error={}",
                    mcp_id, e
                );
                // 继续清理其他资源
            }
        }

        info!(
            "[RAII] All MCP service resources have been cleaned up, and a total of {} services have been cleaned up.",
            count
        );
        Ok(())
    }

    /// 仅移除服务（依赖 RAII 自动清理进程）
    ///
    /// 从 DashMap 中移除服务，触发 RAII 自动清理。
    /// 不会清理路由和缓存，适用于需要快速移除服务的场景。
    pub fn remove_service(&self, mcp_id: &str) -> bool {
        if self.services.remove(mcp_id).is_some() {
            info!(
                "[RAII] The service has been removed and the process will be automatically cleaned up: mcp_id={}",
                mcp_id
            );
            true
        } else {
            debug!("[RAII] Service does not exist: mcp_id={}", mcp_id);
            false
        }
    }

    /// 重置健康检查失败计数
    ///
    /// 使用 get_mut 避免为不存在的服务创建空 entry
    pub fn reset_probe_failures(&self, mcp_id: &str) {
        if let Some(mut entry) = self.services.get_mut(mcp_id) {
            entry.value_mut().status_mut().reset_probe_failures();
        }
    }

    /// 增加健康检查失败计数，返回增加后的值
    ///
    /// 使用 entry API 确保原子性操作
    pub fn increment_probe_failures(&self, mcp_id: &str) -> u32 {
        self.services
            .get_mut(mcp_id)
            .map(|mut entry| entry.value_mut().status_mut().increment_probe_failures())
            .unwrap_or(0)
    }

    /// 清理资源用于重启（保留配置缓存）
    ///
    /// 与 cleanup_resources 不同，此方法不会清理配置缓存，
    /// 允许后续使用缓存的配置重新启动服务。
    pub async fn cleanup_resources_for_restart(&self, mcp_id: &str) -> Result<()> {
        info!(
            "[RAII] Start cleaning up resources for restart: mcp_id={}",
            mcp_id
        );

        // 创建路径以构建要删除的路由路径
        let mcp_sse_router_path = McpRouterPath::new(mcp_id.to_string(), McpProtocol::Sse)
            .map_err(|e| {
                anyhow::anyhow!("Failed to create SSE router path for {}: {}", mcp_id, e)
            })?;
        let base_sse_path = mcp_sse_router_path.base_path;

        let mcp_stream_router_path = McpRouterPath::new(mcp_id.to_string(), McpProtocol::Stream)
            .map_err(|e| {
                anyhow::anyhow!("Failed to create Stream router path for {}: {}", mcp_id, e)
            })?;
        let base_stream_path = mcp_stream_router_path.base_path;

        // 移除相关路由
        DynamicRouterService::delete_route(&base_sse_path);
        DynamicRouterService::delete_route(&base_stream_path);

        // RAII 核心：从 DashMap 移除会触发 McpProcessGuard::drop()
        if self.services.remove(mcp_id).is_some() {
            info!(
                "[RAII] The service has been removed from the map (for restart), McpProcessGuard will automatically cancel the token: mcp_id={}",
                mcp_id
            );
        } else {
            debug!(
                "[RAII] Service does not exist, skip removal: mcp_id={}",
                mcp_id
            );
        }

        // 注意：不清理配置缓存，保留用于重启
        // self.unregister_mcp_config(mcp_id).await; // 不调用

        // 清理健康状态缓存
        GLOBAL_RESTART_TRACKER.clear_health_status(mcp_id);

        info!(
            "[RAII] MCP service resource cleanup completed (for restart): mcp_id={}",
            mcp_id
        );
        Ok(())
    }
}

/// MCP 配置缓存（使用 moka 实现 TTL）
///
/// ## 存储架构说明
///
/// MCP 配置存储在两个位置：
///
/// 1. **McpServiceStatus.mcp_config**（服务状态中）
///    - 存储当前运行服务的配置
///    - 随服务清理而被删除
///    - 用于快速访问当前服务的配置
///
/// 2. **GLOBAL_MCP_CONFIG_CACHE**（全局缓存）
///    - 独立于服务状态存储
///    - 有 TTL（24 小时）
///    - 用于服务重启时恢复配置
///
/// ## 为什么需要两处存储？
///
/// - 服务清理后，McpServiceStatus 被删除，但配置仍在缓存中
/// - 下次请求到来时，可以从缓存恢复配置并重启服务
/// - 实现了服务的自动重启能力
///
/// ## 优先级
///
/// 1. 请求 header 中的配置（最新）
/// 2. 缓存中的配置（兜底）
///
/// ## TTL
///
/// - 24 小时（可配置）
/// - max_capacity: 1000（防止内存溢出）
pub struct McpConfigCache {
    cache: Cache<String, McpConfig>,
}

impl McpConfigCache {
    pub fn new() -> Self {
        Self {
            cache: Cache::builder()
                .time_to_live(Duration::from_secs(24 * 60 * 60)) // 24 小时 TTL
                .max_capacity(1000) // 最多缓存 1000 个配置，防止内存溢出
                .build(),
        }
    }

    pub async fn insert(&self, mcp_id: String, config: McpConfig) {
        self.cache.insert(mcp_id.clone(), config).await;
        info!("MCP configuration cached: {} (TTL: 24h)", mcp_id);
    }

    pub async fn get(&self, mcp_id: &str) -> Option<McpConfig> {
        self.cache.get(mcp_id).await
    }

    pub async fn invalidate(&self, mcp_id: &str) {
        self.cache.invalidate(mcp_id).await;
    }

    #[allow(dead_code)]
    pub fn invalidate_all(&self) {
        self.cache.invalidate_all();
    }
}

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

// 全局配置缓存单例
pub static GLOBAL_MCP_CONFIG_CACHE: Lazy<McpConfigCache> = Lazy::new(McpConfigCache::default);

/// MCP 服务重启追踪器
///
/// 用于防止服务频繁重启导致的无限循环
///
/// ## 重启限制
///
/// - 最小重启间隔：30 秒
/// - 如果服务在 30 秒内被标记为需要重启，将跳过重启
/// - 这防止了服务启动失败时的无限重启循环
///
/// ## 健康状态缓存
///
/// - 缓存后端健康状态，避免频繁检查
/// - 缓存时间：5 秒（可配置）
/// - 用于减少 `is_mcp_server_ready()` 调用频率
pub struct RestartTracker {
    // mcp_id -> 最后重启时间
    last_restart: DashMap<String, Instant>,
    // mcp_id -> (健康状态, 检查时间)
    health_status: DashMap<String, (bool, Instant)>,
    // mcp_id -> 启动锁，防止并发启动同一服务
    startup_locks: DashMap<String, Arc<Mutex<()>>>,
}

impl RestartTracker {
    pub fn new() -> Self {
        Self {
            last_restart: DashMap::new(),
            health_status: DashMap::new(),
            startup_locks: DashMap::new(),
        }
    }

    /// 获取缓存的健康状态
    ///
    /// 如果缓存未过期（5秒内），返回缓存值
    /// 否则返回 None，表示需要重新检查
    pub fn get_cached_health_status(&self, mcp_id: &str) -> Option<bool> {
        let cache_duration = Duration::from_secs(5); // 5 秒缓存
        let now = Instant::now();

        self.health_status.get(mcp_id).and_then(|entry| {
            let (is_healthy, check_time) = *entry.value();
            if now.duration_since(check_time) < cache_duration {
                Some(is_healthy)
            } else {
                None
            }
        })
    }

    /// 更新健康状态缓存
    pub fn update_health_status(&self, mcp_id: &str, is_healthy: bool) {
        self.health_status
            .insert(mcp_id.to_string(), (is_healthy, Instant::now()));
    }

    /// 清除健康状态缓存
    pub fn clear_health_status(&self, mcp_id: &str) {
        self.health_status.remove(mcp_id);
    }

    /// 检查是否可以重启服务
    ///
    /// 返回 true 表示可以重启，false 表示在冷却期内
    ///
    /// 注意：此方法仅检查是否可以重启，不会自动插入时间戳。
    /// 时间戳只在服务成功启动后通过 `record_restart()` 方法记录。
    pub fn can_restart(&self, mcp_id: &str) -> bool {
        let now = Instant::now();
        let min_restart_interval = Duration::from_secs(30); // 30 秒最小重启间隔

        // 只检查，不自动插入时间戳
        if let Some(last_restart) = self.last_restart.get(mcp_id) {
            let elapsed = now.duration_since(*last_restart);
            if elapsed < min_restart_interval {
                warn!(
                    "Service {} is in the cooldown period and is only {} seconds since its last restart. Restart is skipped.",
                    mcp_id,
                    elapsed.as_secs()
                );
                return false;
            }
        }
        // 不在冷却期内，但不自动更新时间戳
        true
    }

    /// 记录服务成功重启
    ///
    /// 此方法应在服务成功启动后调用，用于记录重启时间戳。
    /// 配合 `can_restart()` 使用，避免在服务启动失败时插入时间戳。
    pub fn record_restart(&self, mcp_id: &str) {
        self.last_restart.insert(mcp_id.to_string(), Instant::now());
        info!(
            "The service started successfully and the restart time was recorded: {}",
            mcp_id
        );
    }

    /// 清除重启时间戳
    ///
    /// 当服务启动失败时，可选择调用此方法清除时间戳，
    /// 允许立即重试而不必等待冷却期。
    #[allow(dead_code)]
    pub fn clear_restart(&self, mcp_id: &str) {
        self.last_restart.remove(mcp_id);
        info!("Restart timestamp cleared for service {}", mcp_id);
    }

    /// 尝试获取服务启动锁
    ///
    /// 返回 Some(OwnedMutexGuard) 表示获取成功，可以继续启动服务
    /// 返回 None 表示服务正在启动中，应该跳过本次启动
    ///
    /// # 使用方式
    ///
    /// ```ignore
    /// if let Some(_guard) = GLOBAL_RESTART_TRACKER.try_acquire_startup_lock(&mcp_id) {
    ///     // 获取到锁，可以启动服务
    ///     let result = start_service().await;
    ///     // _guard 在作用域结束时自动释放
    /// } else {
    ///     // 未获取到锁，服务正在启动中
    ///     return Ok(Response::503);
    /// }
    /// ```
    pub fn try_acquire_startup_lock(&self, mcp_id: &str) -> Option<OwnedMutexGuard<()>> {
        // 使用 entry API 确保原子性，避免竞态条件
        let lock = self
            .startup_locks
            .entry(mcp_id.to_string())
            .or_insert_with(|| Arc::new(Mutex::new(())))
            .clone();

        // 尝试获取 owned 锁，锁会一直保持到返回的 guard 被 drop
        match lock.try_lock_owned() {
            Ok(guard) => Some(guard),
            Err(_) => {
                // 锁被占用，服务正在启动中
                debug!("Service {} is starting, skip this startup", mcp_id);
                None
            }
        }
    }

    /// 清理服务启动锁
    ///
    /// 当服务启动完成或失败后，应该清理启动锁以允许后续重试
    /// 注意：正常情况下锁会随 MutexGuard 自动释放，此方法用于异常清理
    #[allow(dead_code)]
    pub fn cleanup_startup_lock(&self, mcp_id: &str) {
        self.startup_locks.remove(mcp_id);
        debug!("Cleaned startup lock for service {}", mcp_id);
    }
}

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

// 全局重启追踪器单例
pub static GLOBAL_RESTART_TRACKER: Lazy<RestartTracker> = Lazy::new(RestartTracker::default);

// 提供一个便捷的函数来获取全局 ProxyHandlerManager
pub fn get_proxy_manager() -> &'static ProxyHandlerManager {
    &GLOBAL_PROXY_MANAGER
}