websocket_server.lua

-- Simplified WebSocket server for diffusion.nvim
-- Merges TCP server, client management, and handshake handling into single component
-- Based on successful claudecode.nvim architecture

local bit = require('bit')

local WebSocketServer = {}

-- WebSocket GUID as defined in RFC 6455
local WEBSOCKET_GUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"

function WebSocketServer:new(config, events)
  local instance = {
    _config = config or {},
    _events = events,
    _logger = require('diffusion.utils.logger'):child("WebSocketServer"),
    _server = nil,
    _port = nil,
    _auth_token = nil,
    _running = false,
    _clients = {},
    _message_handlers = {},
    _stats = {
      connections_total = 0,
      connections_active = 0,
      handshakes_completed = 0,
      handshakes_failed = 0,
      messages_sent = 0,
      messages_received = 0
    }
  }
  
  setmetatable(instance, { __index = self })
  return instance
end

-- Start WebSocket server
function WebSocketServer:start(port, auth_token)
  if self._running then
    return false, "Server already running"
  end
  
  self._port = port
  self._auth_token = auth_token
  
  self._logger:info("Starting WebSocket server", { port = port, has_auth = auth_token ~= nil })
  
  -- Create TCP server
  self._server = vim.loop.new_tcp()
  if not self._server then
    return false, "Failed to create TCP handle"
  end
  
  -- Bind to localhost only for security
  local bind_success, bind_error = self._server:bind("127.0.0.1", port)
  if not bind_success then
    local error_msg = "Failed to bind to 127.0.0.1:" .. port .. " - " .. (bind_error or "unknown error")
    self._logger:error(error_msg)
    self._server:close()
    return false, error_msg
  end
  
  -- Start listening with connection handler
  self._logger:debug("Setting up TCP listen handler")
  
  local listen_success, listen_error = self._server:listen(128, function(err)
    self._logger:debug("TCP listen callback triggered", { error = err })
    
    if err then
      self._logger:error("Accept error", { error = err })
      return
    end
    
    self._logger:debug("Attempting to handle new connection")
    self:_handle_new_connection()
  end)
  
  if not listen_success then
    local error_msg = "Failed to listen on port " .. port .. " - " .. (listen_error or "unknown error")
    self._logger:error(error_msg)
    self._server:close()
    return false, error_msg
  end
  
  self._running = true
  self._logger:info("WebSocket server started successfully", { port = port })
  
  return true, port
end

-- Stop WebSocket server
function WebSocketServer:stop()
  if not self._running then
    return
  end
  
  self._logger:info("Stopping WebSocket server")
  
  -- Disconnect all clients
  for client_id, client in pairs(self._clients) do
    if client.socket and not client.socket:is_closing() then
      client.socket:close()
    end
  end
  self._clients = {}
  
  -- Close server
  if self._server then
    self._server:close()
    self._server = nil
  end
  
  self._running = false
  self._logger:info("WebSocket server stopped")
end

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

-- Send initial file context immediately after WebSocket connection
function WebSocketServer:_send_initial_file_context(client)
  -- Get current file information
  local file_path = vim.api.nvim_buf_get_name(0)

  -- Debug logging to see what file we're getting
  self._logger:info("Getting current file for initial context", {
    id = client.id,
    file_path = file_path,
    is_empty = file_path == "",
    is_nil = file_path == nil,
    buffer_id = vim.api.nvim_get_current_buf()
  })

  -- Skip if no file or special buffer
  if not file_path or file_path == "" or file_path:match("^%w+://") then
    self._logger:debug("Skipping initial file context - no valid file", {
      id = client.id,
      file_path = file_path,
      reason = not file_path and "nil" or file_path == "" and "empty" or "special buffer"
    })
    return
  end

  -- Get cursor position
  local cursor_pos = vim.api.nvim_win_get_cursor(0)
  local line_num = cursor_pos[1]
  local col_num = cursor_pos[2]

  -- Get current selection if any
  local selection_data = self:_get_current_selection()

  -- Create file context notification
  local file_url = "file://" .. file_path
  local notification = {
    jsonrpc = "2.0",
    method = "selection_changed",
    params = {
      text = selection_data.text or "",
      filePath = file_path,
      fileUrl = file_url,
      selection = {
        start = {
          line = (selection_data.start_line or line_num) - 1,
          character = (selection_data.start_col or (col_num + 1)) - 1
        },
        ["end"] = {
          line = (selection_data.end_line or line_num) - 1,
          character = (selection_data.end_col or (col_num + 1)) - 1
        },
        isEmpty = not selection_data.text or selection_data.text == ""
      }
    }
  }

  self._logger:info("Sending initial file context", {
    id = client.id,
    file = file_path,
    has_selection = selection_data.text and selection_data.text ~= "",
    notification = notification
  })

  local success = self:_send_text_frame(client, vim.json.encode(notification))
  if not success then
    self._logger:warn("Failed to send initial file context", { id = client.id })
  else
    self._logger:info("Successfully sent initial file context", {
      id = client.id,
      file_path = file_path,
      file_url = file_url,
      selection_lines = selection_data.text and #selection_data.text:split("\n") or 0
    })
  end
end

-- Get current selection (simple implementation)
function WebSocketServer:_get_current_selection()
  local mode = vim.fn.mode()

  -- If in visual mode, get selection
  if mode:match("[vV\22]") then
    local start_pos = vim.fn.getpos("'<")
    local end_pos = vim.fn.getpos("'>")

    if start_pos and end_pos and start_pos[2] > 0 and end_pos[2] > 0 then
      local lines = vim.api.nvim_buf_get_lines(0, start_pos[2] - 1, end_pos[2], false)
      local text = table.concat(lines, "\n")

      if #lines == 1 then
        text = string.sub(text, start_pos[3], end_pos[3])
      else
        if #lines > 1 then
          lines[1] = string.sub(lines[1], start_pos[3])
          lines[#lines] = string.sub(lines[#lines], 1, end_pos[3])
        end
        text = table.concat(lines, "\n")
      end

      return {
        text = text,
        start_line = start_pos[2],
        start_col = start_pos[3],
        end_line = end_pos[2],
        end_col = end_pos[3]
      }
    end
  end

  -- No selection
  return {
    text = nil,
    start_line = nil,
    start_col = nil,
    end_line = nil,
    end_col = nil
  }
end

-- Check if has connected clients
function WebSocketServer:has_clients()
  return self:get_client_count() > 0
end

-- Get client count
function WebSocketServer:get_client_count()
  return self._stats.connections_active
end

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

-- Set message handler
function WebSocketServer:on_message(handler)
  table.insert(self._message_handlers, handler)
end

-- Set client connected callback
function WebSocketServer:on_client_connected(handler)
  self._on_client_connected = handler
end

-- Set client disconnected callback
function WebSocketServer:on_client_disconnected(handler)
  self._on_client_disconnected = handler
end

-- Install a provider-mutex pair consulted during the handshake.
--   acquire(): atomic claim — returns (true) on success or (false, reason).
--   release(): undo a prior successful acquire.
-- The server invokes acquire() after auth + single-client checks. If the
-- handshake response then fails to send, release() is called immediately
-- so the lock isn't leaked on transport errors. Normal disconnects rely on
-- on_client_disconnected to release.
function WebSocketServer:set_provider_mutex(mutex)
  self._provider_mutex = mutex
end

-- Broadcast message to all connected clients
function WebSocketServer:broadcast_message(message)
  local json_message = vim.json.encode(message)
  local sent_count = 0
  
  for client_id, client in pairs(self._clients) do
    if client.state == "connected" then
      local success = self:_send_text_frame(client, json_message)
      if success then
        sent_count = sent_count + 1
      end
    end
  end
  
  return sent_count > 0
end

-- Send message to specific client
function WebSocketServer:send_message(client_id, message)
  self._logger:debug("Attempting to send message to client", {
    client_id = client_id,
    message_method = message.method,
    message_id = message.id,
    has_result = message.result ~= nil,
    has_error = message.error ~= nil
  })
  
  local client = self._clients[client_id]
  if not client then
    self._logger:error("Client not found for message send", { client_id = client_id })
    return false
  end
  
  if client.state ~= "connected" then
    self._logger:error("Client not connected for message send", { 
      client_id = client_id,
      client_state = client.state
    })
    return false
  end
  
  local json_encode_success, json_message = pcall(vim.json.encode, message)
  if not json_encode_success then
    self._logger:error("Failed to JSON encode message", {
      client_id = client_id,
      error = json_message
    })
    return false
  end
  
  self._logger:debug("Encoded message to JSON", {
    client_id = client_id,
    json_length = #json_message
  })
  
  return self:_send_text_frame(client, json_message)
end

-- Handle new TCP connection
function WebSocketServer:_handle_new_connection()
  self._logger:debug("_handle_new_connection called")
  
  -- Accept connection
  local client_socket = vim.loop.new_tcp()
  if not client_socket then
    self._logger:error("Failed to create client socket")
    return
  end
  self._logger:debug("Created client socket")
  
  local accept_success, accept_error = self._server:accept(client_socket)
  if not accept_success then
    self._logger:error("Failed to accept connection", { error = accept_error })
    client_socket:close()
    return
  end
  self._logger:debug("Accepted client connection")
  
  -- Get client address and validate localhost only
  local peername = client_socket:getpeername()
  self._logger:debug("Got peer name", { peername = peername })
  
  if not peername or (peername.ip ~= "127.0.0.1" and peername.ip ~= "::1") then
    self._logger:warn("Rejected non-localhost connection", { address = peername and peername.ip })
    client_socket:close()
    return
  end
  
  local client_addr = peername.ip .. ":" .. peername.port
  local client_id = require('diffusion.utils').generate_id("ws_client")
  
  -- Create client object
  local client = {
    id = client_id,
    socket = client_socket,
    address = client_addr,
    state = "handshaking",
    buffer = "",
    created_at = os.time(),
    last_ping = os.time()
  }
  
  -- Store client
  self._clients[client_id] = client
  self._stats.connections_total = self._stats.connections_total + 1
  self._stats.connections_active = self._stats.connections_active + 1
  
  self._logger:debug("New connection accepted", { id = client_id, address = client_addr })
  
  -- Set up client handlers
  self:_setup_client_handlers(client)
end

-- Setup handlers for client connection
function WebSocketServer:_setup_client_handlers(client)
  local socket = client.socket
  
  -- Handle incoming data
  socket:read_start(function(err, data)
    if err then
      self._logger:error("Client read error", { id = client.id, error = err })
      self:_disconnect_client(client.id, "read_error")
      return
    end
    
    if data then
      self:_handle_client_data(client, data)
    else
      -- Connection closed
      self:_disconnect_client(client.id, "connection_closed")
    end
  end)
  
  -- Note: vim.loop TCP handles don't have 'on' method
  -- Errors are handled through the read_start callback above
end

-- Handle incoming data from client with optimized string concatenation
function WebSocketServer:_handle_client_data(client, data)
  -- Use table.concat for better performance with frequent concatenations
  if not client.buffer_parts then
    client.buffer_parts = {}
    client.buffer_parts[1] = client.buffer
  end
  table.insert(client.buffer_parts, data)
  
  -- Always keep client.buffer current: _handle_handshake and _handle_websocket_frames
  -- read it directly. Consolidate the parts array only when it grows large to limit
  -- repeated concat cost on high-frequency data.
  client.buffer = table.concat(client.buffer_parts)
  if #client.buffer_parts > 10 then
    client.buffer_parts = {client.buffer}
  end
  
  if client.state == "handshaking" then
    self:_handle_handshake(client)
  elseif client.state == "connected" then
    self:_handle_websocket_frames(client)
  end
end

-- Handle WebSocket handshake
function WebSocketServer:_handle_handshake(client)
  self._logger:debug("Starting WebSocket handshake", { id = client.id, buffer_len = #client.buffer })
  
  -- Check for complete HTTP request
  local header_end = client.buffer:find("\r\n\r\n")
  if not header_end then
    self._logger:debug("Incomplete HTTP request, waiting for more data", { id = client.id })
    return -- Need more data
  end
  
  local request_data = client.buffer:sub(1, header_end + 3)
  client.buffer = client.buffer:sub(header_end + 4)
  
  -- Reset buffer parts after processing handshake
  client.buffer_parts = {client.buffer}
  
  -- Log raw HTTP request for debugging
  self._logger:debug("Raw HTTP request from Claude Code", {
    id = client.id,
    request_preview = request_data:sub(1, 500),
    request_lines = vim.split(request_data, "\r\n")
  })
  
  -- Parse HTTP request
  local request = self:_parse_http_request(request_data)
  if not request then
    self._logger:warn("Invalid HTTP request", { 
      id = client.id,
      raw_request = request_data:sub(1, 200)
    })
    self:_send_error_response(client, 400, "Bad Request")
    self:_disconnect_client(client.id, "invalid_request")
    return
  end
  
  self._logger:debug("Parsed HTTP request headers", {
    id = client.id,
    method = request.method,
    path = request.path,
    version = request.version,
    header_count = vim.tbl_count(request.headers),
    has_subprotocol = request.headers["sec-websocket-protocol"] ~= nil,
    subprotocol = request.headers["sec-websocket-protocol"]
  })
  
  -- Validate WebSocket upgrade
  if not self:_validate_websocket_request(request) then
    self._logger:warn("Invalid WebSocket upgrade", { id = client.id })
    self:_send_error_response(client, 400, "Bad WebSocket Request")
    self:_disconnect_client(client.id, "invalid_websocket")
    return
  end
  
  -- Authenticate request
  self._logger:debug("Authenticating request", { 
    id = client.id, 
    has_auth_token = self._auth_token ~= nil,
    auth_header = request.headers["x-claude-code-ide-authorization"] and "present" or "missing"
  })
  
  if not self:_authenticate_request(request) then
    self._logger:warn("Authentication failed", { id = client.id })
    self._stats.handshakes_failed = self._stats.handshakes_failed + 1
    self:_send_error_response(client, 401, "Unauthorized")
    self:_disconnect_client(client.id, "auth_failed")
    return
  end
  
  self._logger:debug("Authentication successful", { id = client.id })

  -- Enforce single-client policy: reject if another client is already connected
  local existing_client_id = self:_get_connected_client_id()
  if existing_client_id then
    self._logger:warn("Connection rejected - another client already connected", {
      new_client = client.id,
      existing_client = existing_client_id
    })
    self._stats.handshakes_failed = self._stats.handshakes_failed + 1
    self:_send_error_response(client, 409, "Another client is already connected")
    self:_disconnect_client(client.id, "single_client_policy")
    return
  end

  -- Cross-handler provider mutex: rejects when another protocol owns the slot.
  -- Atomic claim — if successful, the lock must be released either via the
  -- on-disconnect path (normal case) or inline below if response send fails.
  local mutex_acquired = false
  if self._provider_mutex then
    local ok, reason = self._provider_mutex.acquire()
    if not ok then
      self._logger:warn("Connection rejected by provider mutex", {
        id = client.id, reason = reason
      })
      self._stats.handshakes_failed = self._stats.handshakes_failed + 1
      self:_send_error_response(client, 503, reason or "Provider locked")
      self:_disconnect_client(client.id, "provider_locked")
      return
    end
    mutex_acquired = true
  end

  -- Generate handshake response
  local response = self:_create_handshake_response(request)
  local success = self:_send_raw_data(client, response)
  if not success then
    if mutex_acquired then
      self._provider_mutex.release()
    end
    self:_disconnect_client(client.id, "handshake_send_failed")
    return
  end
  
  -- Complete handshake
  client.state = "connected"
  self._stats.handshakes_completed = self._stats.handshakes_completed + 1

  self._logger:info("WebSocket handshake completed", { id = client.id })

  -- Call connected callback if set
  if self._on_client_connected then
    vim.schedule(function()
      self._on_client_connected(client)
    end)
  end

  -- CRITICAL: Send current file context immediately after handshake
  -- Use vim.schedule to avoid fast event context error
  vim.schedule(function()
    -- Small delay to ensure Claude is fully initialized
    vim.defer_fn(function()
      self:_send_initial_file_context(client)

      -- ALSO: Trigger UI to send current file context through normal selection system
      -- This ensures Claude gets the file through the proper selection tracking system
      local ui = require('diffusion.ui')
      if ui and ui.get_instance then
        local ui_instance = ui.get_instance()
        if ui_instance then
          if ui_instance.track_selection then
            ui_instance:track_selection()
          end
          -- EXTRA: Force file context send if UI has _send_file_context method
          if ui_instance._send_file_context then
            ui_instance:_send_file_context()
          end
        end
      end

    end, 100) -- 100ms delay to ensure Claude is ready
  end)

  -- Process any remaining data as WebSocket frames
  if client.buffer ~= "" then
    self:_handle_websocket_frames(client)
  end
end

-- Parse HTTP request
function WebSocketServer:_parse_http_request(data)
  local lines = {}
  for line in data:gmatch("[^\r\n]+") do
    table.insert(lines, line)
  end
  
  if #lines == 0 then
    return nil
  end
  
  -- Parse request line
  local method, path, version = lines[1]:match("^(%S+)%s+(%S+)%s+(.+)$")
  if method ~= "GET" then
    return nil
  end
  
  -- Parse headers
  local headers = {}
  for i = 2, #lines do
    local name, value = lines[i]:match("^([^:]+):%s*(.+)$")
    if name and value then
      headers[name:lower()] = value
    end
  end
  
  return {
    method = method,
    path = path,
    version = version,
    headers = headers
  }
end

-- Validate WebSocket upgrade request
function WebSocketServer:_validate_websocket_request(request)
  local headers = request.headers
  
  -- Validate HTTP method and version (following claudecode.nvim pattern)
  if request.method ~= "GET" then
    self._logger:debug("Invalid HTTP method for WebSocket", { method = request.method })
    return false
  end
  
  if not request.version or not request.version:match("^HTTP/1%.1") then
    self._logger:debug("Invalid HTTP version for WebSocket", { version = request.version })
    return false
  end
  
  -- Check required headers
  if not headers.upgrade or headers.upgrade:lower() ~= "websocket" then
    self._logger:debug("Missing or invalid Upgrade header", { upgrade = headers.upgrade })
    return false
  end
  
  if not headers.connection or not headers.connection:lower():find("upgrade") then
    self._logger:debug("Missing or invalid Connection header", { connection = headers.connection })
    return false
  end
  
  if not headers["sec-websocket-key"] then
    self._logger:debug("Missing Sec-WebSocket-Key header")
    return false
  end
  
  -- Validate WebSocket key format (should be base64 encoded 16 bytes = 24 characters)
  local key = headers["sec-websocket-key"]
  if #key ~= 24 then
    self._logger:debug("Invalid Sec-WebSocket-Key format", { 
      key_length = #key,
      expected_length = 24
    })
    return false
  end
  
  if not headers["sec-websocket-version"] or headers["sec-websocket-version"] ~= "13" then
    self._logger:debug("Missing or unsupported Sec-WebSocket-Version", { 
      version = headers["sec-websocket-version"]
    })
    return false
  end
  
  self._logger:debug("WebSocket request validation passed")
  return true
end

-- Authenticate WebSocket request
function WebSocketServer:_authenticate_request(request)
  if not self._auth_token then
    return true -- No authentication required
  end

  local headers = request.headers

  -- Check configured auth headers (primary method)
  local header_names = (self._config and self._config.auth_headers) or {
    "x-claude-code-ide-authorization",  -- Claude Code
    "x-opencode-ide-authorization",     -- OpenCode
    "x-websocket-secret",
    "websocket-secret",
  }
  for _, h in ipairs(header_names) do
    local v = headers[h]
    if v and v == self._auth_token then
      return true
    end
  end

  -- Check Authorization header (fallback)
  local auth_header = headers.authorization
  if auth_header then
    local token = auth_header:match("^Bearer%s+(.+)$")
    if token == self._auth_token then
      return true
    end
  end

  -- Check query-string token (used by clients that cannot send custom headers,
  -- e.g. Node's built-in WebSocket / browser clients).
  -- Caveat: query strings can leak into Referer headers and server access logs.
  -- Safe today because the server is bound to 127.0.0.1; if that binding ever
  -- changes, switch token-bearing clients to the Authorization header.
  if request.path then
    local query_string = request.path:match("%?(.+)$")
    if query_string then
      for param in query_string:gmatch("([^&]+)") do
        local key, value = param:match("^([^=]+)=(.+)$")
        if key == "token" and value == self._auth_token then
          return true
        end
      end
    end
  end

  return false
end

-- Create WebSocket handshake response
function WebSocketServer:_create_handshake_response(request)
  local ws_key = request.headers["sec-websocket-key"]
  local accept_key = self:_generate_accept_key(ws_key)
  
  local response_lines = {
    "HTTP/1.1 101 Switching Protocols",
    "Upgrade: websocket",
    "Connection: Upgrade",
    "Sec-WebSocket-Accept: " .. accept_key,
    "Server: diffusion.nvim/1.0"
  }
  
  -- Handle WebSocket subprotocol if requested
  local requested_protocol = request.headers["sec-websocket-protocol"]
  if requested_protocol then
    self._logger:debug("Client requested WebSocket subprotocol", {
      protocol = requested_protocol
    })
    
    -- Accept the requested protocol (Claude Code may require this)
    table.insert(response_lines, "Sec-WebSocket-Protocol: " .. requested_protocol)
  end
  
  -- Add empty lines to end headers
  table.insert(response_lines, "")
  table.insert(response_lines, "")
  
  return table.concat(response_lines, "\r\n")
end

-- Generate WebSocket accept key
function WebSocketServer:_generate_accept_key(ws_key)
  local crypto = require('diffusion.utils.crypto')
  local concatenated = ws_key .. WEBSOCKET_GUID
  local sha1_hash = crypto.sha1(concatenated)
  return crypto.base64_encode(sha1_hash)
end

-- Handle WebSocket frames
function WebSocketServer:_handle_websocket_frames(client)
  while #client.buffer >= 2 do
    local frame = self:_parse_websocket_frame(client)
    if not frame then
      break -- Need more data or invalid frame
    end
    
    self:_process_websocket_frame(client, frame)
  end
end

-- Parse single WebSocket frame
function WebSocketServer:_parse_websocket_frame(client)
  local buffer = client.buffer
  
  if #buffer < 2 then
    return nil
  end
  
  local byte1 = buffer:byte(1)
  local byte2 = buffer:byte(2)
  
  -- Use math operations instead of bitwise for Lua 5.1 compatibility  
  local fin = math.floor(byte1 / 128) == 1
  local opcode = byte1 % 16
  local masked = math.floor(byte2 / 128) == 1
  local payload_len = byte2 % 128
  
  local header_size = 2
  local mask_size = masked and 4 or 0
  
  -- Determine actual payload length
  local actual_payload_len = payload_len
  if payload_len == 126 then
    if #buffer < header_size + 2 then return nil end
    actual_payload_len = buffer:byte(3) * 256 + buffer:byte(4)
    header_size = header_size + 2
  elseif payload_len == 127 then
    if #buffer < header_size + 8 then return nil end
    -- Use lower 32 bits only
    actual_payload_len = buffer:byte(7) * 16777216 + buffer:byte(8) * 65536 + 
                        buffer:byte(9) * 256 + buffer:byte(10)
    header_size = header_size + 8
  end
  
  -- Check if complete frame available
  local total_size = header_size + mask_size + actual_payload_len
  if #buffer < total_size then
    return nil
  end
  
  -- Extract payload
  local payload_start = header_size + mask_size + 1
  local payload = buffer:sub(payload_start, payload_start + actual_payload_len - 1)
  
  -- Unmask payload if needed (optimized: uses bit.bxor for LuaJIT, chunked for large payloads)
  if masked then
    local mask_key = buffer:sub(header_size + 1, header_size + 4)
    local len = #payload
    if len > 0 then
      -- Pre-extract mask bytes once
      local m1, m2, m3, m4 = mask_key:byte(1, 4)
      local mask = { m1, m2, m3, m4 }

      -- Chunked unmasking to avoid Lua unpack limit (~7998 args in LuaJIT)
      -- Use smaller chunks (1000) for reliable cross-platform behavior
      local CHUNK_SIZE = 1000
      local chunks = {}
      for chunk_start = 1, len, CHUNK_SIZE do
        local chunk_end = math.min(chunk_start + CHUNK_SIZE - 1, len)
        local chunk = {}
        for i = chunk_start, chunk_end do
          chunk[i - chunk_start + 1] = bit.bxor(payload:byte(i), mask[((i - 1) % 4) + 1])
        end
        table.insert(chunks, string.char(unpack(chunk)))
      end
      payload = table.concat(chunks)
    end
  end
  
  -- Remove processed data from buffer
  client.buffer = buffer:sub(total_size + 1)
  
  -- Reset buffer parts after processing frame
  client.buffer_parts = {client.buffer}
  
  return {
    fin = fin,
    opcode = opcode,
    payload = payload
  }
end

-- Process WebSocket frame
function WebSocketServer:_process_websocket_frame(client, frame)
  if frame.opcode == 0x1 then
    -- Text frame - handle as JSON message
    self:_handle_text_message(client, frame.payload)
  elseif frame.opcode == 0x8 then
    -- Close frame
    self:_disconnect_client(client.id, "close_frame")
  elseif frame.opcode == 0x9 then
    -- Ping frame - respond with pong
    self:_send_pong_frame(client, frame.payload)
  elseif frame.opcode == 0xa then
    -- Pong frame - update ping time
    client.last_ping = os.time()
  end
end

-- Handle text message from client
function WebSocketServer:_handle_text_message(client, text)
  self._stats.messages_received = self._stats.messages_received + 1
  
  self._logger:debug("Raw message received", { 
    id = client.id, 
    text_length = #text,
    text_preview = text:sub(1, 200)
  })
  
  -- Parse JSON message
  local success, message = pcall(vim.json.decode, text)
  if not success then
    self._logger:error("Invalid JSON message", { 
      id = client.id, 
      error = message,
      text = text:sub(1, 500)
    })
    return
  end
  
  self._logger:debug("MCP message received", { 
    id = client.id, 
    method = message.method,
    has_id = message.id ~= nil,
    message_type = type(message),
    message_keys = message and vim.tbl_keys(message) or {}
  })
  
  -- Call message handlers
  local handler_count = 0
  for _, handler in ipairs(self._message_handlers) do
    handler_count = handler_count + 1
    local handler_success, handler_error = pcall(handler, client.id, message)
    if not handler_success then
      self._logger:error("Message handler failed", {
        id = client.id,
        handler_index = handler_count,
        error = handler_error
      })
    end
  end
  
  self._logger:debug("Message handlers called", {
    id = client.id,
    handler_count = handler_count
  })
end

-- Send text frame to client
function WebSocketServer:_send_text_frame(client, text)
  self._logger:debug("Sending WebSocket text frame", {
    id = client.id,
    text_length = #text,
    text_preview = text:sub(1, 200)
  })
  
  local frame_data = self:_create_text_frame(text)
  local success = self:_send_raw_data(client, frame_data)
  
  if success then
    self._stats.messages_sent = self._stats.messages_sent + 1
    self._logger:debug("WebSocket text frame sent successfully", { id = client.id })
  else
    self._logger:error("Failed to send WebSocket text frame", { id = client.id })
  end
  
  return success
end

-- Send pong frame
function WebSocketServer:_send_pong_frame(client, payload)
  local frame_data = self:_create_pong_frame(payload or "")
  return self:_send_raw_data(client, frame_data)
end

-- Create text frame
function WebSocketServer:_create_text_frame(text)
  local payload_len = #text
  local header = string.char(0x81) -- FIN=1, opcode=1 (text)
  
  if payload_len < 126 then
    header = header .. string.char(payload_len)
  elseif payload_len < 65536 then
    header = header .. string.char(126) .. 
             string.char(math.floor(payload_len / 256) % 256) .. 
             string.char(payload_len % 256)
  else
    header = header .. string.char(127) ..
             string.rep("\0", 4) .. -- High 32 bits
             string.char(math.floor(payload_len / 16777216) % 256) ..
             string.char(math.floor(payload_len / 65536) % 256) ..
             string.char(math.floor(payload_len / 256) % 256) ..
             string.char(payload_len % 256)
  end
  
  return header .. text
end

-- Create pong frame
function WebSocketServer:_create_pong_frame(payload)
  local payload_len = #payload
  local header = string.char(0x8A) -- FIN=1, opcode=A (pong)
  
  if payload_len < 126 then
    header = header .. string.char(payload_len)
  else
    -- Simplified - assume small pong payloads
    header = header .. string.char(payload_len)
  end
  
  return header .. payload
end

-- Send raw data to client
function WebSocketServer:_send_raw_data(client, data)
  if not client.socket or client.socket:is_closing() then
    return false
  end
  
  local success, error = client.socket:write(data)
  if not success then
    self._logger:error("Failed to send data", { id = client.id, error = error })
    return false
  end
  
  return true
end

-- Send error response
function WebSocketServer:_send_error_response(client, status_code, status_text)
  local body = status_text
  local response_lines = {
    "HTTP/1.1 " .. status_code .. " " .. status_text,
    "Content-Type: text/plain",
    "Content-Length: " .. #body,
    "Connection: close",
    "Server: diffusion.nvim/1.0",
    "",
    body
  }
  
  local response = table.concat(response_lines, "\r\n")
  self:_send_raw_data(client, response)
end

-- Disconnect client
function WebSocketServer:_disconnect_client(client_id, reason)
  local client = self._clients[client_id]
  if not client then
    return
  end

  self._logger:debug("Disconnecting client", { id = client_id, reason = reason })

  -- Call disconnected callback if set and client was connected
  if self._on_client_disconnected and client.state == "connected" then
    vim.schedule(function()
      self._on_client_disconnected(client, reason)
    end)
  end

  -- Release the provider mutex — only previously-connected clients ever held
  -- it. Handshake-fail-after-acquire is released inline at the call site, so
  -- gating on "connected" prevents double-release.
  if self._provider_mutex and client.state == "connected" then
    self._provider_mutex.release()
  end

  -- Close socket
  if client.socket and not client.socket:is_closing() then
    client.socket:close()
  end
  
  -- Update stats
  self._stats.connections_active = math.max(0, self._stats.connections_active - 1)
  
  -- Remove from clients
  self._clients[client_id] = nil
end

-- Get the ID of an already-connected client, or nil if none
-- Used to enforce single-client policy (prevent multiple Claude instances)
function WebSocketServer:_get_connected_client_id()
  for client_id, client in pairs(self._clients) do
    if client.state == "connected" then
      return client_id
    end
  end
  return nil
end

-- Get server statistics
function WebSocketServer:get_stats()
  return vim.deepcopy(self._stats)
end

return WebSocketServer