dotfiles

local encode = function( width, height, cell )
  local rows, styles = {}, {}
  for row = 1, height do
    local text, runs = {}, {}
    for col = 1, width do
      local char, id, attributes = cell( row, col )
      text[#text + 1] = char
      local key = tostring( id )
      styles[key] = styles[key] or attributes
      local previous = runs[#runs]
      if previous and previous.style == key then
        previous.width = previous.width + 1
      else
        runs[#runs + 1] = { col = col, width = 1, style = key, }
      end
    end
    rows[row] = { text = table.concat( text ), runs = runs, }
  end
  return { width = width, height = height, rows = rows, styles = styles, }
end

local consume = function( state, events )
  for _, event in ipairs( events ) do
    local name = event[1]
    if name == 'flush' then state.ready = true end
    for i = 2, #event do
      local item = event[i]
      if name == 'hl_attr_define' then
        local attributes = item[2]
        attributes.cterm = item[3]
        state.styles[item[1]] = attributes
      elseif name == 'default_colors_set' then
        state.defaults = {
          foreground = item[1], background = item[2], special = item[3],
          ctermfg = item[4], ctermbg = item[5],
        }
      elseif name == 'grid_resize' and item[1] == 1 then
        state.width, state.height = item[2], item[3]
        state.cells = {}
        for row = 1, state.height do
          state.cells[row] = {}
          for col = 1, state.width do state.cells[row][col] = { ' ', 0, } end
        end
      elseif name == 'grid_clear' and item[1] == 1 then
        for row = 1, state.height do
          for col = 1, state.width do state.cells[row][col] = { ' ', 0, } end
        end
      elseif name == 'grid_line' and item[1] == 1 then
        local row, col, id = item[2] + 1, item[3] + 1, 0
        for _, cell in ipairs( item[4] ) do
          id = cell[2] or id
          for _ = 1, cell[3] or 1 do
            state.cells[row][col] = { cell[1], id, }
            col = col + 1
          end
        end
      elseif name == 'grid_scroll' and item[1] == 1 then
        local before = vim.deepcopy( state.cells )
        for row = item[2] + 1, item[3] do
          for col = item[4] + 1, item[5] do
            local source_row, source_col = row + item[6], col + item[7]
            state.cells[row][col] = source_row > item[2] and source_row <= item[3]
                and source_col > item[4] and source_col <= item[5]
                and before[source_row][source_col] or { ' ', 0, }
          end
        end
      elseif name == 'grid_cursor_goto' and item[1] == 1 then
        state.cursor = { row = item[2] + 1, col = item[3] + 1, }
      end
    end
  end
end

local decode = function( state, unpacker, chunk )
  local offset = 1
  while offset <= #chunk do
    local message
    message, offset = unpacker( chunk, offset )
    if message then
      if message[1] == 2 and message[2] == 'redraw' then
        consume( state, message[3] )
      elseif message[1] == 1 and message[2] == 1 and message[3] ~= vim.NIL then
        state.error = vim.inspect( message[3] )
      end
    end
  end
end

local info = function()
  local uis = vim.api.nvim_list_uis()
  if #uis == 0 then error( 'nvim_capture requires an attached UI; headless buffer text is not a rendered grid' ) end
  local grid = { width = vim.o.columns, height = vim.o.lines, }
  grid.coordinate_base = 1
  grid.mode = vim.api.nvim_get_mode().mode
  grid.uis = uis
  grid.windows = {}
  for _, window in ipairs( vim.api.nvim_tabpage_list_wins( 0 ) ) do
    local position = vim.api.nvim_win_get_position( window )
    grid.windows[#grid.windows + 1] = {
      window = window, buffer = vim.api.nvim_win_get_buf( window ),
      row = position[1] + 1, col = position[2] + 1,
      width = vim.api.nvim_win_get_width( window ), height = vim.api.nvim_win_get_height( window ),
      config = vim.api.nvim_win_get_config( window ),
    }
  end
  grid.external = {}
  for _, name in ipairs( { 'ext_cmdline', 'ext_messages', 'ext_popupmenu', 'ext_tabline', 'ext_wildmenu', } ) do
    for _, ui in ipairs( uis ) do
      if ui[name] then grid.external[#grid.external + 1] = name; break end
    end
  end
  return grid
end

local capture = function( address, metadata )
  local state = { cells = {}, styles = { [0] = { cterm = {}, }, }, }
  local unpacker = vim.mpack.Unpacker()
  local host, port = address:match( '^([%w%.%-]+):(%d+)$' )
  local channel = host and vim.uv.new_tcp() or vim.uv.new_pipe( false )
  local read = function( problem, chunk )
    if problem then state.error = tostring( problem ); return end
    if not chunk then
      if not state.ready then state.error = 'UI capture connection closed before flush' end
      return
    end
    local ok, failure = pcall( decode, state, unpacker, chunk )
    if not ok then state.error = tostring( failure ) end
  end
  local connected = function( problem )
    if problem then state.error = tostring( problem ); return end
    channel:read_start( read )
    channel:write( vim.mpack.encode( { 0, 1, 'nvim_ui_attach', {
      metadata.width, metadata.height,
      { rgb = true, ext_linegrid = true, ext_multigrid = false, },
    }, } ) )
  end
  local ok, result = pcall( function()
    if host then channel:connect( host, tonumber( port ), connected ) else channel:connect( address, connected ) end
    local ready = vim.wait( 5000, function() return state.ready or state.error ~= nil end, 10 )
    if state.error then error( state.error ) end
    if not ready then error( 'timed out waiting for a rendered UI flush' ) end
    local grid = encode( state.width, state.height, function( row, col )
      local cell = state.cells[row][col]
      return cell[1], cell[2], state.styles[cell[2]]
    end )
    grid.coordinate_base = metadata.coordinate_base
    grid.defaults = state.defaults
    grid.cursor = state.cursor
    grid.cursor.mode = metadata.mode
    grid.uis = metadata.uis
    grid.windows = metadata.windows
    grid.external = metadata.external
    return grid
  end )
  channel:read_stop()
  channel:close()
  if not ok then error( result ) end
  return result
end

return { info = info, capture = capture, encode = encode, consume = consume, decode = decode, }