codex.lua

-- Codex protocol handler for diffusion.nvim
-- Implements MCP over WebSocket using the same tool/message flow as Claude, but with Codex discovery and tagging

local CodexHandler = {}

-- Shared MCP modules (tools registry, message dispatcher, selection broadcaster)
local tools = require('diffusion.protocol.mcp.tools')
local messages = require('diffusion.protocol.mcp.messages')
local selection = require('diffusion.protocol.mcp.selection')
local diff_events = require('diffusion.protocol.mcp.diff_events')
local codex_diff = require('diffusion.protocol.codex.diff')

function CodexHandler:new(config, events_bus, diff_manager, provider_lock)
  local instance = {
    _config = config,
    _events = events_bus,
    _logger = require('diffusion.utils.logger'):child("Codex"),
    _server = nil,
    _port = nil,
    _auth_token = nil,
    _running = false,
    _tools = nil,
    _diff_manager = diff_manager,
    _provider_lock = provider_lock,  -- Cross-handler "first connection wins" mutex
    _discovery = require('diffusion.utils.discovery'):new('codex'),
    _event_subscriptions = {},  -- Selection MCP module writes here
    _event_unsubs = {},         -- Shared diff_events helper writes unsub fns here
    _stats = {
      selections_sent = 0,
      tools_called = 0,
      messages_received = 0,
      diffs_accepted = 0,
      diffs_rejected = 0,
    }
  }
  setmetatable(instance, { __index = self })
  return instance
end

function CodexHandler:is_available()
  -- Prefer discovery info; otherwise treat as available (we can start our own server)
  local info = self._discovery:get_service_info('codex') or {}
  return info.available == true or true
end

function CodexHandler:start()
  if self._running then
    self._logger:warn("Codex handler already running", { port = self._port })
    return false, "Already running"
  end

  self._logger:info("Starting Codex handler")

  -- Clean stale codex lockfiles
  pcall(function() self._discovery:cleanup_stale_lock_files() end)

  -- Choose port and token
  local utils = require('diffusion.utils')
  self._port = utils.find_available_port(10000, 65535)
  if not self._port then
    return false, "No available port for Codex"
  end

  local crypto = require('diffusion.utils.crypto')
  self._auth_token = crypto.generate_secure_token(true)

  -- Start local WS server (token-protected; echoes requested subprotocol like 'mcp')
  local WebSocketServer = require('diffusion.server.websocket_server')
  self._server = WebSocketServer:new({}, self._events)

  if self._provider_lock then
    self._server:set_provider_mutex({
      acquire = function() return self._provider_lock:try_acquire("codex") end,
      release = function() self._provider_lock:release("codex") end,
    })
  end

  local ok, err = self._server:start(self._port, self._auth_token)
  if not ok then
    self._server = nil
    self._logger:error("Failed to start Codex WS server", { error = err })
    return false, err
  end

  -- Create Codex lockfile so the stdio bridge (spawned by Codex) can discover us
  local folders = { vim.fn.getcwd() }
  local data = {
    authToken = self._auth_token,
    workspaceFolders = folders,
    ideName = "diffusion.nvim",
    transport = "ws"
  }
  local created = self._discovery:create_lock_file(self._port, data, 'codex')
  if not created then
    self._logger:warn("Failed to create Codex lockfile")
  end

  -- Also expose via env as a fallback for the bridge
  vim.env.DIFFUSION_WEBSOCKET_PORT = tostring(self._port)
  vim.env.DIFFUSION_WEBSOCKET_SECRET = self._auth_token

  -- Initialize tools and handlers
  self:_init_tools()
  self:_setup_message_handlers()
  self:_setup_selection_notifications()
  diff_events.setup(self, { protocol = "codex" })

  self._running = true
  self._logger:info("Codex handler started", { port = self._port, pid = vim.fn.getpid() })
  return true, self._port
end

function CodexHandler:stop()
  if not self._running then return true end

  self._logger:info("Stopping Codex handler")
  for _, sub in ipairs(self._event_subscriptions) do
    self._events:off(sub.event, sub.handler)
  end
  self._event_subscriptions = {}
  for _, unsub in ipairs(self._event_unsubs) do
    pcall(unsub)
  end
  self._event_unsubs = {}
  if self._server then self._server:stop() end
  if self._discovery then self._discovery:cleanup() end
  self._running = false
  return true
end

function CodexHandler:is_connected()
  return self._server and self._server:is_running() and self._server:has_clients()
end

-- MCP selections
function CodexHandler:send_selection(selection_data)
  if not self._server or not self._server:is_running() then
    self._logger:error("WebSocket server not running (codex)")
    return false
  end
  local mcp_message = self:_translate_selection_to_mcp(selection_data)
  self._stats.selections_sent = self._stats.selections_sent + 1
  return self._server:broadcast_message(mcp_message)
end

function CodexHandler:send_message(message)
  -- No tmux coupling for Codex; this is a no-op passthrough.
  self._logger:debug("send_message (codex): no-op", { message = message })
  return true
end

-- Diff entry points
function CodexHandler:open_diff(diff_params)
  diff_params.protocol = "codex"
  return self._diff_manager:show_diff(diff_params)
end

-- Tool initialization (reuse Claude’s tool registry but route to our handlers)
function CodexHandler:_init_tools()
  tools._init_tools(self)
end

-- Tool wrappers
function CodexHandler:_handle_open_file(args)
  return tools._handle_open_file(self, args)
end

function CodexHandler:_handle_open_diff(args)
  return codex_diff._handle_open_diff(self, args)
end

function CodexHandler:_handle_open_diff_deferred(request_info)
  return codex_diff._handle_open_diff_deferred(self, request_info)
end

function CodexHandler:_handle_show_diff(args)
  return codex_diff._handle_show_diff(self, args)
end

function CodexHandler:_handle_dismiss_diff(args)
  return codex_diff._handle_dismiss_diff(self, args)
end

function CodexHandler:_get_pending_diffs()
  return codex_diff._get_pending_diffs(self)
end

-- Message + events setup (reuse Claude modules)
function CodexHandler:_setup_message_handlers()
  return messages._setup_message_handlers(self)
end

function CodexHandler:_setup_selection_notifications()
  return selection._setup_selection_notifications(self)
end

-- Selection serialization
function CodexHandler:_translate_selection_to_mcp(selection_data)
  return {
    jsonrpc = "2.0",
    method = "selection_changed",
    params = {
      text = selection_data.text,
      filePath = selection_data.file_path,
      fileUrl = "file://" .. selection_data.file_path,
      selection = {
        start = { line = selection_data.start_line, character = selection_data.start_col },
        ["end"] = { line = selection_data.end_line, character = selection_data.end_col },
        isEmpty = selection_data.text == ""
      }
    }
  }
end

return CodexHandler