preview.nvim/lua/preview/compiler.lua

local M = {}

local diagnostic = require('preview.diagnostic')
local log = require('preview.log')

---@type table<integer, preview.Process>
local active = {}

---@type table<integer, integer>
local watching = {}

---@type table<integer, true>
local opened = {}

---@type table<integer, string>
local last_output = {}

---@type table<integer, table>
local viewer_procs = {}

---@type table<integer, uv.uv_fs_event_t>
local open_watchers = {}

local debounce_timers = {}

local DEBOUNCE_MS = 500

---@param bufnr integer
local function stop_open_watcher(bufnr)
  local w = open_watchers[bufnr]
  if w then
    w:stop()
    w:close()
    open_watchers[bufnr] = nil
  end
end

---@param bufnr integer
local function close_viewer(bufnr)
  local obj = viewer_procs[bufnr]
  if obj then
    local kill = obj.kill
    kill(obj, 'sigterm')
    viewer_procs[bufnr] = nil
  end
end

---@param bufnr integer
---@param output_file string
---@param open_config boolean|string[]
local function do_open(bufnr, output_file, open_config)
  if open_config == true then
    vim.ui.open(output_file)
  elseif type(open_config) == 'table' then
    local open_cmd = vim.list_extend({}, open_config)
    table.insert(open_cmd, output_file)
    viewer_procs[bufnr] = vim.system(open_cmd)
  end
end

---@param val string[]|fun(ctx: preview.Context): string[]
---@param ctx preview.Context
---@return string[]
local function eval_list(val, ctx)
  if type(val) == 'function' then
    return val(ctx)
  end
  return val
end

---@param val string|fun(ctx: preview.Context): string
---@param ctx preview.Context
---@return string
local function eval_string(val, ctx)
  if type(val) == 'function' then
    return val(ctx)
  end
  return val
end

---@param provider preview.ProviderConfig
---@param ctx preview.Context
---@return string[]?
local function resolve_reload_cmd(provider, ctx)
  if type(provider.reload) == 'function' then
    return provider.reload(ctx)
  elseif type(provider.reload) == 'table' then
    return vim.list_extend({}, provider.reload)
  end
  return nil
end

---@param bufnr integer
---@param name string
---@param provider preview.ProviderConfig
---@param ctx preview.Context
function M.compile(bufnr, name, provider, ctx, opts)
  opts = opts or {}

  if vim.fn.executable(provider.cmd[1]) ~= 1 then
    vim.notify(
      '[preview.nvim]: "' .. provider.cmd[1] .. '" is not executable (run :checkhealth preview)',
      vim.log.levels.ERROR
    )
    return
  end

  if vim.bo[bufnr].modified then
    vim.cmd('silent! update')
  end

  if active[bufnr] then
    log.dbg('killing existing process for buffer %d before recompile', bufnr)
    M.stop(bufnr)
  end

  local output_file = ''
  if provider.output then
    output_file = eval_string(provider.output, ctx)
  end

  local resolved_ctx = vim.tbl_extend('force', ctx, { output = output_file })

  local cwd = ctx.root
  if provider.cwd then
    cwd = eval_string(provider.cwd, resolved_ctx)
  end

  if output_file ~= '' then
    last_output[bufnr] = output_file
  end

  local reload_cmd
  if not opts.oneshot then
    reload_cmd = resolve_reload_cmd(provider, resolved_ctx)
  end

  if reload_cmd then
    log.dbg(
      'starting long-running process for buffer %d with provider "%s": %s',
      bufnr,
      name,
      table.concat(reload_cmd, ' ')
    )

    local stderr_acc = {}
    local obj
    obj = vim.system(
      reload_cmd,
      {
        cwd = cwd,
        env = provider.env,
        stderr = vim.schedule_wrap(function(_err, data)
          if not data or not vim.api.nvim_buf_is_valid(bufnr) then
            return
          end
          stderr_acc[#stderr_acc + 1] = data
          local errors_mode = provider.errors
          if errors_mode == nil then
            errors_mode = 'diagnostic'
          end
          if provider.error_parser and errors_mode then
            local output = table.concat(stderr_acc)
            if errors_mode == 'diagnostic' then
              diagnostic.set(bufnr, name, provider.error_parser, output, ctx)
            elseif errors_mode == 'quickfix' then
              local ok, diags = pcall(provider.error_parser, output, ctx)
              if ok and diags and #diags > 0 then
                local items = {}
                for _, d in ipairs(diags) do
                  table.insert(items, {
                    bufnr = bufnr,
                    lnum = d.lnum + 1,
                    col = d.col + 1,
                    text = d.message,
                    type = d.severity == vim.diagnostic.severity.WARN and 'W' or 'E',
                  })
                end
                vim.fn.setqflist(items, 'r')
                local win = vim.fn.win_getid()
                vim.cmd.cwindow()
                vim.fn.win_gotoid(win)
              end
            end
          end
        end),
      },
      vim.schedule_wrap(function(result)
        if active[bufnr] and active[bufnr].obj == obj then
          active[bufnr] = nil
        end
        if not vim.api.nvim_buf_is_valid(bufnr) then
          return
        end

        if result.code ~= 0 then
          log.dbg('long-running process failed for buffer %d (exit code %d)', bufnr, result.code)
          local errors_mode = provider.errors
          if errors_mode == nil then
            errors_mode = 'diagnostic'
          end
          if provider.error_parser and errors_mode then
            local output = (result.stdout or '') .. (result.stderr or '')
            if errors_mode == 'diagnostic' then
              diagnostic.set(bufnr, name, provider.error_parser, output, ctx)
            elseif errors_mode == 'quickfix' then
              local ok, diagnostics = pcall(provider.error_parser, output, ctx)
              if ok and diagnostics and #diagnostics > 0 then
                local items = {}
                for _, d in ipairs(diagnostics) do
                  table.insert(items, {
                    bufnr = bufnr,
                    lnum = d.lnum + 1,
                    col = d.col + 1,
                    text = d.message,
                    type = d.severity == vim.diagnostic.severity.WARN and 'W' or 'E',
                  })
                end
                vim.fn.setqflist(items, 'r')
                local win = vim.fn.win_getid()
                vim.cmd.cwindow()
                vim.fn.win_gotoid(win)
              end
            end
          end
          vim.api.nvim_exec_autocmds('User', {
            pattern = 'PreviewCompileFailed',
            data = {
              bufnr = bufnr,
              provider = name,
              code = result.code,
              stderr = result.stderr or '',
            },
          })
        end
      end)
    )

    if provider.open and not opts.oneshot and not opened[bufnr] and output_file ~= '' then
      local pre_stat = vim.uv.fs_stat(output_file)
      local pre_mtime = pre_stat and pre_stat.mtime.sec or 0
      local out_dir = vim.fn.fnamemodify(output_file, ':h')
      local out_name = vim.fn.fnamemodify(output_file, ':t')
      stop_open_watcher(bufnr)
      local watcher = vim.uv.new_fs_event()
      if watcher then
        open_watchers[bufnr] = watcher
        watcher:start(
          out_dir,
          {},
          vim.schedule_wrap(function(err, filename, _events)
            if err or vim.fn.fnamemodify(filename or '', ':t') ~= out_name then
              return
            end
            if opened[bufnr] then
              stop_open_watcher(bufnr)
              return
            end
            if not vim.api.nvim_buf_is_valid(bufnr) then
              stop_open_watcher(bufnr)
              return
            end
            local new_stat = vim.uv.fs_stat(output_file)
            if not (new_stat and new_stat.mtime.sec > pre_mtime) then
              return
            end
            stop_open_watcher(bufnr)
            stderr_acc = {}
            local errors_mode = provider.errors
            if errors_mode == nil then
              errors_mode = 'diagnostic'
            end
            if errors_mode == 'diagnostic' then
              diagnostic.clear(bufnr)
            elseif errors_mode == 'quickfix' then
              vim.fn.setqflist({}, 'r')
              vim.cmd.cwindow()
            end
            do_open(bufnr, output_file, provider.open)
            opened[bufnr] = true
          end)
        )
      end
    end

    active[bufnr] = { obj = obj, provider = name, output_file = output_file, is_reload = true }

    vim.api.nvim_create_autocmd('BufUnload', {
      buffer = bufnr,
      once = true,
      callback = function()
        M.stop(bufnr)
        stop_open_watcher(bufnr)
        if not provider.detach then
          close_viewer(bufnr)
        end
        last_output[bufnr] = nil
      end,
    })

    vim.api.nvim_exec_autocmds('User', {
      pattern = 'PreviewCompileStarted',
      data = { bufnr = bufnr, provider = name },
    })
    return
  end

  local cmd = vim.list_extend({}, provider.cmd)
  if provider.args then
    vim.list_extend(cmd, eval_list(provider.args, resolved_ctx))
  end

  log.dbg('compiling buffer %d with provider "%s": %s', bufnr, name, table.concat(cmd, ' '))

  local obj
  obj = vim.system(
    cmd,
    {
      cwd = cwd,
      env = provider.env,
    },
    vim.schedule_wrap(function(result)
      if active[bufnr] and active[bufnr].obj == obj then
        active[bufnr] = nil
      end
      if not vim.api.nvim_buf_is_valid(bufnr) then
        return
      end

      local errors_mode = provider.errors
      if errors_mode == nil then
        errors_mode = 'diagnostic'
      end

      if result.code == 0 then
        log.dbg('compilation succeeded for buffer %d', bufnr)
        if errors_mode == 'diagnostic' then
          diagnostic.clear(bufnr)
        elseif errors_mode == 'quickfix' then
          vim.fn.setqflist({}, 'r')
          vim.cmd.cwindow()
        end
        vim.api.nvim_exec_autocmds('User', {
          pattern = 'PreviewCompileSuccess',
          data = { bufnr = bufnr, provider = name, output = output_file },
        })
        if provider.reload == true and output_file:match('%.html$') then
          local r = require('preview.reload')
          r.start()
          r.inject(output_file)
          r.broadcast()
        end
        if
          provider.open
          and not opts.oneshot
          and not opened[bufnr]
          and output_file ~= ''
          and vim.uv.fs_stat(output_file)
        then
          do_open(bufnr, output_file, provider.open)
          opened[bufnr] = true
        end
      else
        log.dbg('compilation failed for buffer %d (exit code %d)', bufnr, result.code)
        if provider.error_parser and errors_mode then
          local output = (result.stdout or '') .. (result.stderr or '')
          if errors_mode == 'diagnostic' then
            diagnostic.set(bufnr, name, provider.error_parser, output, ctx)
          elseif errors_mode == 'quickfix' then
            local ok, diagnostics = pcall(provider.error_parser, output, ctx)
            if ok and diagnostics and #diagnostics > 0 then
              local items = {}
              for _, d in ipairs(diagnostics) do
                table.insert(items, {
                  bufnr = bufnr,
                  lnum = d.lnum + 1,
                  col = d.col + 1,
                  text = d.message,
                  type = d.severity == vim.diagnostic.severity.WARN and 'W' or 'E',
                })
              end
              vim.fn.setqflist(items, 'r')
              local win = vim.fn.win_getid()
              vim.cmd.cwindow()
              vim.fn.win_gotoid(win)
            end
          end
        end
        vim.api.nvim_exec_autocmds('User', {
          pattern = 'PreviewCompileFailed',
          data = {
            bufnr = bufnr,
            provider = name,
            code = result.code,
            stderr = result.stderr or '',
          },
        })
      end
    end)
  )

  active[bufnr] = { obj = obj, provider = name, output_file = output_file }

  vim.api.nvim_create_autocmd('BufUnload', {
    buffer = bufnr,
    once = true,
    callback = function()
      M.stop(bufnr)
      if not provider.detach then
        close_viewer(bufnr)
      end
      last_output[bufnr] = nil
    end,
  })

  vim.api.nvim_exec_autocmds('User', {
    pattern = 'PreviewCompileStarted',
    data = { bufnr = bufnr, provider = name },
  })
end

---@param bufnr integer
function M.stop(bufnr)
  local proc = active[bufnr]
  if not proc then
    return
  end
  log.dbg('stopping process for buffer %d', bufnr)
  ---@type fun(self: table, signal: string|integer)
  local kill = proc.obj.kill
  kill(proc.obj, 'sigterm')

  local timer = vim.uv.new_timer()
  if timer then
    timer:start(5000, 0, function()
      timer:close()
      if active[bufnr] and active[bufnr].obj == proc.obj then
        kill(proc.obj, 'sigkill')
        active[bufnr] = nil
      end
    end)
  end
end

function M.stop_all()
  for bufnr, _ in pairs(active) do
    M.stop(bufnr)
  end
  for bufnr, _ in pairs(watching) do
    M.unwatch(bufnr)
  end
  for bufnr, _ in pairs(open_watchers) do
    stop_open_watcher(bufnr)
  end
  for bufnr, _ in pairs(viewer_procs) do
    close_viewer(bufnr)
  end
  require('preview.reload').stop()
end

---@param bufnr integer
---@param name string
---@param provider preview.ProviderConfig
---@param ctx_builder fun(bufnr: integer): preview.Context
function M.toggle(bufnr, name, provider, ctx_builder)
  local is_longrunning = type(provider.reload) == 'table' or type(provider.reload) == 'function'

  if is_longrunning then
    if active[bufnr] then
      M.stop(bufnr)
      vim.notify('[preview.nvim]: watching stopped', vim.log.levels.INFO)
    else
      M.compile(bufnr, name, provider, ctx_builder(bufnr))
      vim.notify('[preview.nvim]: watching with "' .. name .. '"', vim.log.levels.INFO)
    end
    return
  end

  if watching[bufnr] then
    M.unwatch(bufnr)
    vim.notify('[preview.nvim]: watching stopped', vim.log.levels.INFO)
    return
  end

  local au_id = vim.api.nvim_create_autocmd('BufWritePost', {
    buffer = bufnr,
    callback = function()
      if debounce_timers[bufnr] then
        debounce_timers[bufnr]:stop()
      else
        debounce_timers[bufnr] = vim.uv.new_timer()
      end
      debounce_timers[bufnr]:start(
        DEBOUNCE_MS,
        0,
        vim.schedule_wrap(function()
          local ctx = ctx_builder(bufnr)
          M.compile(bufnr, name, provider, ctx)
        end)
      )
    end,
  })

  watching[bufnr] = au_id
  log.dbg('watching buffer %d with provider "%s"', bufnr, name)
  vim.notify('[preview.nvim]: watching with "' .. name .. '"', vim.log.levels.INFO)

  vim.api.nvim_create_autocmd('BufUnload', {
    buffer = bufnr,
    once = true,
    callback = function()
      M.unwatch(bufnr)
      stop_open_watcher(bufnr)
      if not provider.detach then
        close_viewer(bufnr)
      end
      opened[bufnr] = nil
    end,
  })

  M.compile(bufnr, name, provider, ctx_builder(bufnr))
end

---@param bufnr integer
function M.unwatch(bufnr)
  local au_id = watching[bufnr]
  if not au_id then
    return
  end
  vim.api.nvim_del_autocmd(au_id)
  if debounce_timers[bufnr] then
    debounce_timers[bufnr]:stop()
    debounce_timers[bufnr]:close()
    debounce_timers[bufnr] = nil
  end
  watching[bufnr] = nil
  log.dbg('unwatched buffer %d', bufnr)
end

---@param bufnr integer
---@param name string
---@param provider preview.ProviderConfig
---@param ctx preview.Context
function M.clean(bufnr, name, provider, ctx)
  if not provider.clean then
    vim.notify(
      '[preview.nvim]: provider "' .. name .. '" has no clean command',
      vim.log.levels.WARN
    )
    return
  end

  local output_file = ''
  if provider.output then
    output_file = eval_string(provider.output, ctx)
  end
  local resolved_ctx = vim.tbl_extend('force', ctx, { output = output_file })

  local cmd = eval_list(provider.clean, resolved_ctx)
  local cwd = resolved_ctx.root
  if provider.cwd then
    cwd = eval_string(provider.cwd, resolved_ctx)
  end

  log.dbg('cleaning buffer %d with provider "%s": %s', bufnr, name, table.concat(cmd, ' '))

  vim.system(
    cmd,
    { cwd = cwd },
    vim.schedule_wrap(function(result)
      if result.code == 0 then
        log.dbg('clean succeeded for buffer %d', bufnr)
        vim.notify('[preview.nvim]: clean complete', vim.log.levels.INFO)
      else
        log.dbg('clean failed for buffer %d (exit code %d)', bufnr, result.code)
        vim.notify('[preview.nvim]: clean failed: ' .. (result.stderr or ''), vim.log.levels.ERROR)
      end
    end)
  )
end

---@param bufnr integer
---@return boolean
function M.open(bufnr, open_config)
  local output = last_output[bufnr]
  if not output then
    log.dbg('no last output file for buffer %d', bufnr)
    return false
  end
  if not vim.uv.fs_stat(output) then
    log.dbg('output file no longer exists for buffer %d: %s', bufnr, output)
    return false
  end
  do_open(bufnr, output, open_config)
  return true
end

---@param bufnr integer
---@return preview.Status
function M.status(bufnr)
  local proc = active[bufnr]
  if proc then
    return {
      compiling = not proc.is_reload,
      watching = watching[bufnr] ~= nil or proc.is_reload == true,
      provider = proc.provider,
      output_file = proc.output_file,
    }
  end
  return { compiling = false, watching = watching[bufnr] ~= nil }
end

M._test = {
  active = active,
  watching = watching,
  opened = opened,
  last_output = last_output,
  debounce_timers = debounce_timers,
  viewer_procs = viewer_procs,
  open_watchers = open_watchers,
}

return M