port_mapping.lua

local log = require('opencode.log')
local config = require('opencode.config')
local util = require('opencode.util')
local OpencodeServer = require('opencode.opencode_server')

local M = {}

-- Signal 0 only checks if a process exists, doesn't actually signal it
local SIG_PID_EXISTS = 0

--- @class PortMappingEntry
--- @field pid number
--- @field directory string
--- @field mode string

--- @class PortMapping
--- @field directory string
--- @field nvim_pids PortMappingEntry[]
--- @field auto_kill boolean
--- @field started_by_nvim boolean
--- @field url string|nil The URL the opencode server is listening on
--- @field server_pid number|nil The PID of the opencode server process (local servers only)

--- @return string
local function file_path()
  return vim.fn.stdpath('data') .. '/opencode_port_mappings.json'
end

--- @return table<string, PortMapping>
local function load()
  local file = io.open(file_path(), 'r')
  if not file then
    return {}
  end
  local content = file:read('*all')
  file:close()
  local ok, data = pcall(vim.json.decode, content or '')
  return ok and data or {}
end

--- @param mappings table<string, PortMapping>
local function save(mappings)
  local path = file_path()
  local file = io.open(path, 'w')
  if not file then
    log.warn('port_mapping: could not open %s for writing', path)
    return
  end
  file:write(vim.json.encode(mappings))
  file:close()
end

--- @param entry PortMappingEntry
--- @return boolean
local function pid_alive(entry)
  return vim.fn.getpid() == entry.pid or vim.uv.kill(entry.pid, SIG_PID_EXISTS) == 0
end

--- Fire-and-forget graceful shutdown request to a server with no clients.
--- Also force-kills the process if server_pid is available.
--- @param port number
--- @param server_pid number|nil
local function kill_orphaned_server(port, server_pid)
  local server_url = config.server.url or '127.0.0.1'
  local normalized_url = util.normalize_url_protocol(server_url)
  local base_url = string.format('%s:%d', normalized_url, port)

  log.info('port_mapping: sending shutdown to orphaned server at %s (server_pid=%s)', base_url, tostring(server_pid))

  OpencodeServer.request_graceful_shutdown(base_url)

  if server_pid then
    OpencodeServer.kill_pid(server_pid)
  else
    log.debug('port_mapping: no server PID available, relying on graceful shutdown only')
  end
end

--- Purge dead nvim PIDs from every mapping and kill any newly-orphaned servers.
local function clean_stale()
  local mappings = load()
  local changed = false

  for port_key, mapping in pairs(mappings) do
    mapping.nvim_pids = mapping.nvim_pids or {}
    local before = #mapping.nvim_pids

    mapping.nvim_pids = vim.tbl_filter(pid_alive, mapping.nvim_pids)

    if #mapping.nvim_pids < before then
      changed = true
    end

    if #mapping.nvim_pids == 0 then
      local port = tonumber(port_key)
      if port and mapping.started_by_nvim then
        kill_orphaned_server(port, mapping.server_pid)
      end
      log.debug('port_mapping: removing port %s (no connected clients)', port_key)
      mappings[port_key] = nil
    end
  end

  if changed then
    save(mappings)
  end
end

--- Return the directory a port is already mapped to, or nil when the port is
--- either free or already mapped to current_dir.
--- @param port number
--- @param current_dir string
--- @return string|nil
function M.mapped_directory(port, current_dir)
  clean_stale()
  local mapping = load()[tostring(port)]
  if mapping and mapping.directory and mapping.directory ~= current_dir then
    return mapping.directory
  end
end

--- Return an existing port serving current_dir, or nil.
--- @param current_dir string
--- @return number|nil
function M.find_port_for_directory(current_dir)
  clean_stale()
  for port_key, mapping in pairs(load()) do
    if mapping.directory == current_dir and mapping.nvim_pids and #mapping.nvim_pids > 0 then
      local port = tonumber(port_key)
      if port then
        return port
      end
    end
  end
end

--- Record that this nvim instance is using the given port.
--- @param port number
--- @param directory string
--- @param started_by_nvim boolean
--- @param mode? string 'serve'|'attach'|'custom'
--- @param url? string The URL the server is listening on
--- @param server_pid? number The PID of the server process (local servers only)
function M.register(port, directory, started_by_nvim, mode, url, server_pid)
  mode = mode or 'serve'
  clean_stale()

  local mappings = load()
  local port_key = tostring(port)
  local current_pid = vim.fn.getpid()
  local auto_kill = config.server.auto_kill

  if not mappings[port_key] then
    mappings[port_key] = {
      directory = directory,
      nvim_pids = {},
      auto_kill = auto_kill,
      started_by_nvim = started_by_nvim,
    }
  end

  local mapping = mappings[port_key]
  mapping.nvim_pids = mapping.nvim_pids or {}
  if mapping.auto_kill == nil then
    mapping.auto_kill = auto_kill
  end
  if mapping.started_by_nvim == nil then
    mapping.started_by_nvim = started_by_nvim
  end
  if url then
    mapping.url = url
  end
  -- Only update server_pid if provided (don't overwrite existing PID with nil)
  if server_pid then
    mapping.server_pid = server_pid
  end

  local pid_exists = false
  local updated = {}
  for _, entry in ipairs(mapping.nvim_pids) do
    table.insert(updated, entry)
    if entry.pid == current_pid then
      pid_exists = true
    end
  end
  mapping.nvim_pids = updated

  if not pid_exists then
    table.insert(mapping.nvim_pids, { pid = current_pid, directory = directory, mode = mode })
  end

  save(mappings)
  log.debug(
    'port_mapping.register: port=%d dir=%s pid=%d mode=%s started_by_nvim=%s auto_kill=%s url=%s server_pid=%s',
    port,
    directory,
    current_pid,
    mode,
    tostring(started_by_nvim),
    tostring(auto_kill),
    tostring(url),
    tostring(server_pid)
  )
end

--- Remove this nvim instance from a port's client list.
--- Shuts the server down when it was the last client and auto_kill is set.
--- Also shuts down attach-mode processes unconditionally.
--- @param port number|nil
--- @param server OpencodeServer instance (state.opencode_server)
function M.unregister(port, server)
  if not port then
    return
  end

  clean_stale()
  local mappings = load()
  local port_key = tostring(port)
  local mapping = mappings[port_key]
  if not mapping then
    return
  end

  local current_pid = vim.fn.getpid()
  local remaining = {}
  for _, entry in ipairs(mapping.nvim_pids or {}) do
    if entry.pid ~= current_pid then
      table.insert(remaining, entry)
    end
  end
  mapping.nvim_pids = remaining

  local should_shutdown = #remaining == 0 and mapping.started_by_nvim and mapping.auto_kill

  if server then
    local is_last_client = #remaining == 0 and mapping.started_by_nvim
    if server.mode == 'attach' then
      if is_last_client then
        log.debug('port_mapping.unregister: last attached client for port %d, killing server', port)
        if mapping.server_pid then
          kill_orphaned_server(port, mapping.server_pid)
        end
      end
    elseif is_last_client then
      local auto_kill_custom_server = config.server.auto_kill and config.server.kill_command
      local server_is_owned = server.job
      log.debug(
        'port_mapping.unregister: last nvim instance for port %d, killing orphaned server',
        port,
        tostring(server_is_owned),
        tostring(auto_kill_custom_server)
      )
      if auto_kill_custom_server or server_is_owned then
        server:shutdown()
      end
    end
  elseif should_shutdown then
    log.debug('port_mapping.unregister: no server object, killing orphaned server for port %d', port)
    kill_orphaned_server(port, mapping.server_pid)
  end

  if should_shutdown then
    mappings[port_key] = nil
  else
    log.debug('port_mapping.unregister: port=%d still has %d client(s)', port, #remaining)
  end

  save(mappings)
end

--- Return the started_by_nvim flag for a port, or false if unknown.
--- @param port number
--- @return boolean
function M.started_by_nvim(port)
  local mapping = load()[tostring(port)]
  return mapping and mapping.started_by_nvim or false
end

--- Find any existing server port (regardless of directory)
--- @return number|nil port number if found, nil otherwise
function M.find_any_existing_port()
  clean_stale()
  local mappings = load()

  for port_key, mapping in pairs(mappings) do
    if mapping.nvim_pids and #mapping.nvim_pids > 0 then
      local port = tonumber(port_key)
      if port then
        return port
      end
    end
  end

  return nil
end

return M