init.lua

-- WebSocket server manager for diffusion.nvim
-- Manages server lifecycle, connections, and protocol routing

local ServerManager = {}

function ServerManager:new(config, events)
  local instance = {
    _config = config,
    _events = events,
    _logger = require('diffusion.utils.logger'):child("Server"),
    _tcp_server = nil,
    _client_manager = nil,
    _port = nil,
    _auth_token = nil,
    _running = false,
    _stats = {
      connections_accepted = 0,
      connections_rejected = 0,
      messages_processed = 0,
      uptime_start = nil
    }
  }
  
  setmetatable(instance, { __index = self })
  return instance
end

-- Start the WebSocket server
function ServerManager:start()
  if self._running then
    self._logger:warn("Server already running", { port = self._port })
    return true, self._port
  end
  
  self._logger:info("Starting WebSocket server")
  
  -- Find available port
  local utils = require('diffusion.utils')
  self._port = utils.find_available_port(
    self._config.port_range[1], 
    self._config.port_range[2]
  )
  
  if not self._port then
    local error_msg = "No available ports in range " .. 
                      self._config.port_range[1] .. "-" .. self._config.port_range[2]
    self._logger:error(error_msg)
    return false, error_msg
  end
  
  -- Generate authentication token
  local crypto = require('diffusion.utils.crypto')
  self._auth_token = crypto.generate_secure_token(true) -- Include timestamp
  
  -- Initialize server components
  self:_init_components()
  
  -- Start TCP server
  local tcp_success, tcp_error = self._tcp_server:start(self._port)
  if not tcp_success then
    self._logger:error("Failed to start TCP server", { error = tcp_error })
    return false, tcp_error
  end
  
  self._running = true
  self._stats.uptime_start = os.time()
  
  self._logger:info("WebSocket server started successfully", {
    port = self._port,
    protocol = self._config.protocol,
    auth_token = self._auth_token
  })
  
  -- Emit server started event
  self._events:emit("server:started", {
    port = self._port,
    protocol = self._config.protocol
  })
  
  return true, self._port
end

-- Stop the WebSocket server
function ServerManager:stop()
  if not self._running then
    return
  end
  
  self._logger:info("Stopping WebSocket server")
  
  -- Stop client manager
  if self._client_manager then
    self._client_manager:disconnect_all()
  end
  
  -- Stop TCP server
  if self._tcp_server then
    self._tcp_server:stop()
  end
  
  self._running = false
  
  self._logger:info("WebSocket server stopped", {
    uptime_seconds = os.time() - (self._stats.uptime_start or os.time())
  })
  
  -- Emit server stopped event
  self._events:emit("server:stopped", {
    port = self._port,
    protocol = self._config.protocol
  })
end

-- Check if server is running
function ServerManager:is_running()
  return self._running
end

-- Get server port
function ServerManager:get_port()
  return self._port
end

-- Get authentication token
function ServerManager:get_auth_token()
  return self._auth_token
end

-- Check if server has connected clients
function ServerManager:has_clients()
  if not self._client_manager then
    return false
  end
  
  return self._client_manager:get_client_count() > 0
end

-- Broadcast message to all connected clients
function ServerManager:broadcast_message(message)
  if not self._running or not self._client_manager then
    self._logger:warn("Cannot broadcast message - server not running")
    return false
  end
  
  local success = self._client_manager:broadcast_message(message)
  if success then
    self._stats.messages_processed = self._stats.messages_processed + 1
  end
  
  return success
end

-- Send message to specific client
function ServerManager:send_message(client_id, message)
  if not self._running or not self._client_manager then
    self._logger:warn("Cannot send message - server not running")
    return false
  end
  
  local success = self._client_manager:send_message(client_id, message)
  if success then
    self._stats.messages_processed = self._stats.messages_processed + 1
  end
  
  return success
end

-- Get server statistics
function ServerManager:get_stats()
  local stats = vim.deepcopy(self._stats)
  stats.running = self._running
  stats.port = self._port
  stats.protocol = self._config.protocol
  stats.uptime_seconds = self._stats.uptime_start and 
                         (os.time() - self._stats.uptime_start) or 0
  
  if self._client_manager then
    stats.active_clients = self._client_manager:get_client_count()
    stats.client_stats = self._client_manager:get_stats()
  else
    stats.active_clients = 0
    stats.client_stats = {}
  end
  
  return stats
end

-- Initialize server components
function ServerManager:_init_components()
  -- Initialize TCP server
  local TcpServer = require('diffusion.server.tcp')
  self._tcp_server = TcpServer:new({
    host = self._config.host,
    logger = self._logger:child("TCP"),
    events = self._events
  })
  
  -- Initialize client manager
  local ClientManager = require('diffusion.server.client')
  self._client_manager = ClientManager:new({
    auth_token = self._auth_token,
    protocol = self._config.protocol,
    max_connections = self._config.max_connections,
    keepalive_interval = self._config.keepalive_interval,
    logger = self._logger:child("Client"),
    events = self._events
  })
  
  -- Set up connection handling
  self._tcp_server:on_connection(function(client_socket, client_addr)
    self._stats.connections_accepted = self._stats.connections_accepted + 1
    self._client_manager:handle_new_connection(client_socket, client_addr)
  end)
  
  -- Set up connection rejection handling
  self._tcp_server:on_connection_rejected(function(reason)
    self._stats.connections_rejected = self._stats.connections_rejected + 1
    self._logger:warn("Connection rejected", { reason = reason })
  end)
  
  -- Set up message handling
  self._client_manager:on_message(function(client_id, message)
    self:_handle_client_message(client_id, message)
  end)
  
  -- Set up client disconnect handling
  self._client_manager:on_disconnect(function(client_id, reason)
    self._logger:debug("Client disconnected", { 
      client_id = client_id, 
      reason = reason 
    })
    
    self._events:emit("client:disconnected", {
      client_id = client_id,
      reason = reason,
      protocol = self._config.protocol
    })
  end)
end

-- Handle messages from clients
function ServerManager:_handle_client_message(client_id, message)
  self._logger:debug("Message received from client", {
    client_id = client_id,
    method = message.method,
    has_id = message.id ~= nil
  })
  
  -- Emit message event for protocol handlers
  self._events:emit("websocket:message_received", {
    client_id = client_id,
    protocol = self._config.protocol,
    message = message
  })
  
  -- Handle different message types
  if message.method == "tools/call" then
    self:_handle_tool_call(client_id, message)
  elseif message.method == "tools/list" then
    self:_handle_tool_list(client_id, message)
  elseif message.method == "initialize" then
    self:_handle_initialize(client_id, message)
  elseif message.method == "notifications/initialized" then
    -- Handle initialized notification (no response needed)
    self._logger:debug("Client initialized notification received")
  else
    self._logger:warn("Unknown message method", {
      client_id = client_id,
      method = message.method
    })
  end
end

-- Handle tool call messages
function ServerManager:_handle_tool_call(client_id, message)
  local tool_name = message.params and message.params.name
  local tool_args = message.params and message.params.arguments or {}
  
  if not tool_name then
    self:_send_error_response(client_id, message.id, "Missing tool name")
    return
  end
  
  self._logger:info("Tool call received", {
    client_id = client_id,
    tool = tool_name,
    id = message.id
  })
  
  -- Tools are now handled directly by protocol handlers
  local success = false
  local result = { error = "Tools are now handled directly by protocol handlers" }
  
  -- Send error response
  do
    local response = {
      jsonrpc = "2.0",
      id = message.id
    }
    
    if success then
      response.result = result
    else
      response.error = {
        code = -1,
        message = tostring(result)
      }
    end
    
    self._client_manager:send_message(client_id, response)
  end
end

-- Handle tool list requests
function ServerManager:_handle_tool_list(client_id, message)
  -- Tools are now handled directly by protocol handlers
  local response = {
    jsonrpc = "2.0",
    id = message.id,
    result = {
      tools = {} -- Empty - tools handled by protocol handlers
    }
  }
  
  self._client_manager:send_message(client_id, response)
end

-- Handle initialization messages
function ServerManager:_handle_initialize(client_id, message)
  local response = {
    jsonrpc = "2.0",
    id = message.id,
    result = {
      protocolVersion = "2025-03-26",
      capabilities = {
        tools = {}
      },
      serverInfo = {
        name = "diffusion.nvim",
        version = "1.0.0",
        protocol = self._config.protocol
      }
    }
  }
  
  self._client_manager:send_message(client_id, response)
  
  self._logger:info("Client initialized", {
    client_id = client_id,
    protocol = self._config.protocol
  })
end

-- Send error response to client
function ServerManager:_send_error_response(client_id, request_id, error_message)
  local response = {
    jsonrpc = "2.0",
    id = request_id,
    error = {
      code = -1,
      message = error_message
    }
  }
  
  self._client_manager:send_message(client_id, response)
end

-- Cleanup method for shutdown
function ServerManager:cleanup()
  self:stop()
end

return ServerManager