opsxshell/src/stores/connection/poolStore.ts

/**
 * Pool Store - 连接池核心管理
 * 负责: 连接 Map、LRU 淘汰、引用计数、创建/断开连接
 */

import { create } from 'zustand'
import type {
  PooledConnection,
  ConnectionConfig,
  ConnectionState,
  PoolClientConfig,
} from './types'
import type { CreateConnectionOptions, CreateConnectionResult, PoolStats } from './types'
import { useSettingsStore } from '../settingsStore'
import { connectionReuseFingerprint } from '../../utils/connectionFingerprint'

// ============== Pool State ==============

interface PoolState {
  // 连接池: connectionId -> PooledConnection
  connections: Map<string, PooledConnection>

  // 配置
  config: PoolClientConfig

  // Actions
  setConfig: (config: Partial<PoolClientConfig>) => void

  // 连接管理
  createConnection: (
    tabId: string,
    config: ConnectionConfig,
    options?: CreateConnectionOptions
  ) => Promise<CreateConnectionResult>

  acquireConnection: (
    tabId: string,
    sessionId: string,
    getSessionConfig: () => Promise<ConnectionConfig | null>
  ) => Promise<{ success: boolean; connectionId?: string; error?: string }>

  resumeConnection: (connectionId: string) => Promise<boolean>

  disconnectConnection: (connectionId: string, reason?: string) => Promise<void>

  // The native service has already torn this connection down. Do not issue a
  // second disconnect command, but release every renderer-side resource.
  forgetNativeConnection: (connectionId: string) => void

  releaseConnection: (connectionId: string, tabId: string) => Promise<void>

  // 连接查询
  getConnection: (connectionId: string) => PooledConnection | undefined
  getConnectionState: (connectionId: string) => ConnectionState | undefined

  // LRU 淘汰
  evictOldestBackgroundConnection: () => Promise<string | null>
  shouldEvict: () => boolean

  // 统计
  getStats: () => PoolStats

  // 断开所有
  disconnectAll: () => Promise<void>
}

// ============== Pool Store Implementation ==============

export const usePoolStore = create<PoolState>((set, get) => ({
  connections: new Map(),
  config: {
    maxConnections: 10,
    maxBackgroundConnections: 5,
    idleTimeoutMs: 5 * 60 * 1000,
    suspendTimeoutMs: 30 * 60 * 1000,
    keepaliveIntervalMs: 60 * 1000,
    maxReconnectAttempts: 3,
  },

  setConfig: (newConfig) => {
    set((state) => ({
      config: { ...state.config, ...newConfig },
    }))
  },

  createConnection: async (tabId, config, _options = {}): Promise<CreateConnectionResult> => {
    const { connections, config: poolConfig } = get()
    const now = Date.now()
    const reuseKey = connectionReuseFingerprint(config)

    // 1. 检查是否已存在相同配置的连接 (连接复用)
    for (const [existingId, conn] of connections.entries()) {
      if (
        connectionReuseFingerprint(conn.config) === reuseKey &&
        conn.state !== 'disconnected' &&
        conn.state !== 'error'
      ) {
        // 复用已有连接
        conn.refCount++
        conn.tabIds.add(tabId)
        conn.activeTabId = tabId

        if (conn.state === 'suspended') {
          // 需要恢复连接 - 简化处理：断开后重建
          const resumed = await get().resumeConnection(existingId)
          if (!resumed) {
            conn.refCount--
            conn.tabIds.delete(tabId)
            continue // 尝试创建新连接
          }
        } else {
          conn.state = 'connected'
        }

        set((state) => {
          const newConnections = new Map(state.connections)
          newConnections.set(existingId, conn)
          return { connections: newConnections }
        })

        return { success: true, connectionId: existingId, isExisting: true }
      }
    }

    // 2. 检查是否需要淘汰
    if (connections.size >= poolConfig.maxConnections) {
      const evictedId = await get().evictOldestBackgroundConnection()
      if (!evictedId) {
        return {
          success: false,
          error: `连接数达到上限 (${poolConfig.maxConnections})，请先关闭一些连接`,
        }
      }
    }

    // 3. 创建新连接
    console.log('>>> Creating new SSH connection via opsxApi...')
    try {
      const sshSettings = useSettingsStore.getState().settings
      const result = await window.opsxApi?.sshConnect({
        host: config.host,
        port: config.port || 22,
        username: config.username,
        password: config.password,
        privateKey: config.privateKey,
        passphrase: config.passphrase,
        proxy: config.proxy,
        jumpHost: config.jumpHost,
        jumpHosts: config.jumpHosts,
        timeoutSeconds: config.timeoutSeconds ?? sshSettings.timeout,
        keepAlive: config.keepAlive ?? sshSettings.keepAlive,
      })
      console.log('>>> sshConnect result:', result)

      if (!result?.success || !result.connectionId) {
        console.log('>>> sshConnect FAILED:', result?.error)
        return {
          success: false,
          error: result?.error || 'Failed to create connection',
        }
      }

      const connectionId = result.connectionId
      console.log('>>> Connection created, connectionId:', connectionId)

      // 设置保活定时器
      const keepaliveInterval = (config.keepAlive ?? sshSettings.keepAlive)
        ? window.setInterval(() => {
            const conn = get().connections.get(connectionId)
            if (conn && (conn.state === 'connected' || conn.state === 'background')) {
              void window.opsxApi?.sshKeepalive(connectionId)
                .then((result) => {
                  if (!result.success) {
                    void get().disconnectConnection(connectionId, result.error || 'keepalive failed')
                  }
                })
                .catch(() => {
                  void get().disconnectConnection(connectionId, 'keepalive failed')
                })
            }
          }, poolConfig.keepaliveIntervalMs)
        : undefined

      // 4. 创建连接池项
      const pooledConnection: PooledConnection = {
        connectionId,
        config: {
          host: config.host,
          port: config.port || 22,
          username: config.username,
          password: config.password,
          privateKey: config.privateKey,
          passphrase: config.passphrase,
          proxy: config.proxy,
          jumpHost: config.jumpHost,
          jumpHosts: config.jumpHosts,
          timeoutSeconds: config.timeoutSeconds ?? sshSettings.timeout,
          keepAlive: config.keepAlive ?? sshSettings.keepAlive,
          startupCommand: config.startupCommand,
        },
        state: 'connected',
        refCount: 1,
        tabIds: new Set([tabId]),
        activeTabId: tabId,
        shellIds: new Set(),
        sftpIds: new Set(),
        stats: {
          connectedAt: now,
          lastActivityAt: now,
          bytesTransferred: 0,
          commandsExecuted: 0,
        },
        keepaliveInterval,
        reconnectAttempts: 0,
        maxReconnectAttempts: poolConfig.maxReconnectAttempts,
      }

      set((state) => {
        const newConnections = new Map(state.connections)
        newConnections.set(connectionId, pooledConnection)
        return { connections: newConnections }
      })

      return { success: true, connectionId }
    } catch (error: unknown) {
      return {
        success: false,
        error: error instanceof Error ? error.message : 'Unknown error creating connection',
      }
    }
  },

  // 辅助: 恢复连接 (简化实现)
  resumeConnection: async (connectionId: string): Promise<boolean> => {
    const conn = get().connections.get(connectionId)
    if (!conn) return false

    if (conn.state === 'connected' || conn.state === 'background') {
      return true
    }

    try {
      const result = await window.opsxApi?.sshConnect(conn.config)
      if (!result?.success) return false

      const newConnectionId = result.connectionId
      if (!newConnectionId) return false

      set((state) => {
        const newConnections = new Map(state.connections)
        newConnections.delete(connectionId)

        const now = Date.now()
        const keepaliveInterval = conn.config.keepAlive === false
          ? undefined
          : window.setInterval(() => {
              void window.opsxApi?.sshKeepalive(newConnectionId).catch(() => {})
            }, state.config.keepaliveIntervalMs)

        newConnections.set(newConnectionId, {
          ...conn,
          connectionId: newConnectionId,
          state: 'connected',
          stats: {
            ...conn.stats,
            lastActivityAt: now,
          },
          keepaliveInterval,
          idleTimeout: undefined,
          reconnectAttempts: 0,
        })

        return { connections: newConnections }
      })

      return true
    } catch {
      return false
    }
  },

  acquireConnection: async (tabId, _sessionId, getSessionConfigFn) => {
    const { connections } = get()

    // 检查是否已有该标签的连接
    for (const [connId, conn] of connections.entries()) {
      if (conn.tabIds.has(tabId) && conn.state !== 'disconnected') {
        return { success: true, connectionId: connId }
      }
    }

    // 获取 session 配置并创建连接
    const config = await getSessionConfigFn()
    if (!config) {
      return { success: false, error: '无法获取会话配置' }
    }

    return get().createConnection(tabId, config)
  },

  disconnectConnection: async (connectionId, _reason = 'manual') => {
    const conn = get().connections.get(connectionId)
    if (!conn) return

    // 清除定时器
    if (conn.keepaliveInterval) {
      clearInterval(conn.keepaliveInterval)
    }
    if (conn.idleTimeout) {
      clearTimeout(conn.idleTimeout)
    }

    try {
      await window.opsxApi?.sshDisconnect(connectionId)
    } catch {
      // 忽略断开错误
    }

    set((state) => {
      const newConnections = new Map(state.connections)
      newConnections.delete(connectionId)
      return { connections: newConnections }
    })
  },

  forgetNativeConnection: (connectionId) => {
    const conn = get().connections.get(connectionId)
    if (!conn) return

    if (conn.keepaliveInterval) clearInterval(conn.keepaliveInterval)
    if (conn.idleTimeout) clearTimeout(conn.idleTimeout)

    set((state) => {
      const newConnections = new Map(state.connections)
      newConnections.delete(connectionId)
      return { connections: newConnections }
    })
  },

  releaseConnection: async (connectionId, tabId) => {
    const conn = get().connections.get(connectionId)
    if (!conn) return

    // Idempotent: ignore tabs that already released this connection
    if (!conn.tabIds.has(tabId)) return

    conn.tabIds.delete(tabId)
    conn.refCount = Math.max(0, conn.refCount - 1)

    if (conn.activeTabId === tabId) {
      conn.activeTabId = null
    }

    // 如果没有标签使用，断开连接
    if (conn.refCount === 0 || conn.tabIds.size === 0) {
      await get().disconnectConnection(connectionId, 'no refs')
      return
    }

    set((state) => {
      const newConnections = new Map(state.connections)
      newConnections.set(connectionId, conn)
      return { connections: newConnections }
    })
  },

  getConnection: (connectionId) => {
    return get().connections.get(connectionId)
  },

  getConnectionState: (connectionId) => {
    return get().connections.get(connectionId)?.state
  },

  evictOldestBackgroundConnection: async () => {
    const { connections } = get()

    let oldestId: string | null = null
    let oldestTime = Infinity

    for (const [id, conn] of connections.entries()) {
      if (conn.state === 'background' || conn.state === 'suspended') {
        const lastActivity = conn.stats.lastBackgroundAt || conn.stats.lastActivityAt
        if (lastActivity < oldestTime) {
          oldestTime = lastActivity
          oldestId = id
        }
      }
    }

    if (oldestId) {
      await get().disconnectConnection(oldestId, 'evicted')
    }

    return oldestId
  },

  shouldEvict: () => {
    const { connections, config } = get()
    return connections.size >= config.maxConnections
  },

  getStats: (): PoolStats => {
    const { connections } = get()
    let connected = 0, background = 0, suspended = 0, connecting = 0, error = 0

    connections.forEach((conn) => {
      switch (conn.state) {
        case 'connected': connected++; break
        case 'background': background++; break
        case 'suspended': suspended++; break
        case 'connecting': connecting++; break
        case 'error': error++; break
      }
    })

    return { total: connections.size, connected, background, suspended, connecting, error }
  },

  disconnectAll: async () => {
    const { connections } = get()
    const promises: Promise<void>[] = []

    connections.forEach((conn) => {
      if (conn.keepaliveInterval) {
        clearInterval(conn.keepaliveInterval)
      }
      if (conn.idleTimeout) {
        clearTimeout(conn.idleTimeout)
      }

      promises.push(
        window.opsxApi?.sshDisconnect(conn.connectionId)
          .then(() => {})
          .catch(() => {}) as Promise<void>
      )
    })

    await Promise.all(promises)
    set({ connections: new Map() })
  },
}))