-- based on the ideas from:
-- https://github.com/RRethy/nvim-align

local M = {}

function M.align_lines(pat, line1, line2, preview_ns, preview_bufnr)
  local re = vim.regex(pat)
  local bufnr = vim.api.nvim_get_current_buf()
  local lines = vim.api.nvim_buf_get_lines(bufnr, line1 - 1, line2, false)
  local newlines = {}
  local preview_buf_line = 0

  -- find the longest match
  local max = -1
  for _, line in pairs(lines) do
    local s = re:match_str(line)
    if s and max < s then max = s end
  end

  -- exit if nothing was found
  if max == -1 then return error('nothing found') end

  for i, line in pairs(lines) do
    local s = re:match_str(line)
    if s then
      local rep = max - s
      local changeset = {
        string.sub(line, 1, s),
        string.rep(' ', rep),
        string.sub(line, s+1),
      }
      local newline = table.concat(changeset)

      -- append to changse if not inside the preview callback
      if not preview_ns then
        newlines[#newlines+1] = newline
      end

      if preview_ns ~= nil then
        -- set extmarks inside the live buffer
        vim.api.nvim_buf_set_extmark(bufnr, preview_ns, line1+i - 2, 0,
          {
            hl_mode = 'combine',
            virt_text_pos = 'overlay',
            virt_text = {
              {changeset[1]},
              {changeset[2], 'Substitute'},
              {changeset[3]},
            },
          })

        -- modify preview buffer
        if preview_bufnr ~= nil then
          local prefix = string.format('|%d| ', line1 + i - 1)
          vim.api.nvim_buf_set_lines(preview_bufnr, preview_buf_line, preview_buf_line, false, { prefix .. newline })
          vim.api.nvim_buf_add_highlight(preview_bufnr, preview_ns, 'Substitute', preview_buf_line, #prefix + s, #prefix + s + #changeset[2])
          preview_buf_line = preview_buf_line + 1
        end
      end
    end
  end

  -- only change buffer when not previewing
  if not preview_ns then
    -- exit if nothing was changed
    if #newlines == 0 then return error('nothing changed') end
    vim.api.nvim_buf_set_lines(bufnr, line1 - 1, line2, 0, newlines)
    return 0
  end

  if preview_ns ~= nil then
    -- open preview buffer only if there is more than a single line of change
    return (#preview_buf_line > 1) and 2 or 1
  end
end

function M.align(pat)
    local top, bot = vim.fn.getpos("'<"), vim.fn.getpos("'>")
    M.align_lines(pat, top[2]-1, bot[2])
    vim.fn.setpos("'<", top)
    vim.fn.setpos("'>", bot)
end

local function aligncmd(opts, preview_ns, preview_bufnr)
  return M.align_lines(opts.fargs[1], opts.line1, opts.line2, preview_ns, preview_bufnr)
end

local default_opts = {
  bindings = true
}

function M.setup(opts)
  opts = vim.tbl_extend('keep', opts or {}, default_opts)

  vim.api.nvim_create_user_command('SimpleAlign', aligncmd,
    {nargs = 1, range = '%', preview = aligncmd})

  if opts.bindings then
    vim.keymap.set('v', '<enter>', ':SimpleAlign ')
  end
end

return M