dep/lua/dep.lua

local logger = require('dep.log')
local proc = require('dep.proc')

---all functions for convenience
---@type table
local M = {}

---@type boolean
local initialized

---root package
---@type table
local root

---performance logging
---@type table
local perf = {}

---table of every package where their id is the index
---@type table
local packages = {}

---path to root of where plugins are downloaded
---@type string
local base_dir

--- get execution time of a function
---@param name string name of performance output
---@param code function function to run
---@vararg any arguments for code
function M.benchmark(name, code, ...)
  local start = os.clock()
  code(...)
  perf[name] = os.clock() - start
end

--- get name of package
---@param id string id of the package
---@return string id
---@nodiscard
function M.getpkgname(id)
  local name = id:match("^[%w-_.]+/([%w-_.]+)$")
  if name then
    return name
  else
    error(string.format(
      'invalid name "%s"; must be in the format "user/package"', id))
  end
end

--- tell the parent it has a child and the child it has a parent
---@param parent table parent package
---@param child table child package
function M.link_dependency(parent, child)
  if not parent.dependents[child.id] then
    parent.dependents[child.id] = child
    parent.dependents[#parent.dependents + 1] = child
  end

  if not child.dependencies[parent.id] then
    child.dependencies[parent.id] = parent
    child.dependencies[#child.dependencies + 1] = parent
  end
end

--- register a new package according to the spec
---@param spec table|string the package spec to register
---@param overrides table|nil a package spec that is used to override options
---@return table package a table containing the package spec
---@nodiscard
function M.registerpkg(spec, overrides)
  overrides = overrides or {}

  if type(spec) ~= "table" then
    spec = { spec }
  end

  local id = spec[1]
  local package = packages[id]

  -- if package hasn't been registered already, get the inital spec regisitered
  if not package then
    package = {
      id = id, -- id of the package
      enabled = true, -- whether it's going to be used
      exists = false, -- if the package exists on the filesystem
      lazy = false, -- if the package is lazy loaded in any way
      added = false, -- if the package has been added in vim
      configured = false, -- if the package has been configured
      loaded = false, -- if a package has been loaded
      subtree_loaded = false,
      on_config = {}, -- table of functions to run on config
      on_setup = {}, -- table of function to run on setup
      on_load = {}, -- table of functions to run on load
      lazy_load = {}, -- table of functions to run which will tell the package
                      -- when to load
      dependencies = {}, -- this package's requirements
      dependents = {}, -- packages that require this package
      perf = {}
    }

    packages[id] = package
  end

  -- register the rest of the package spec
  package.name = spec.as or package.name or M.getpkgname(id)
  package.url = spec.url or package.url or ("https://github.com/"..id..".git")
  package.branch = spec.branch or package.branch
  package.dir = base_dir..package.name
  package.commit = spec.commit
  package.pin = overrides.pin or spec.pin or package.pin
  package.enabled = not overrides.disable and not spec.disable and package.enabled
  package.lazy = spec.lazy or package.lazy

  -- make sure that the package exists
  package.exists = vim.fn.isdirectory(package.dir) ~= 0
  package.configured = package.exists

  -- register all the callback functions
  table.insert(package.on_config, spec.config)
  table.insert(package.on_load, spec.load)
  table.insert(package.lazy_load, spec.lazy)

  -- if the current package isn't the root package then it depends on the root
  -- package
  if root and package ~= root then
    M.link_dependency(root, package)
  end

  -- link the dependencies
  if spec.requires then
    if type(spec.requires) == "string" then
      spec.requires = { spec.requires }
    end
    for _, v in pairs(spec.requires) do
      M.link_dependency(M.registerpkg(v), package)
    end
  end

  -- and link the dependents
  if spec.deps then
    if type(spec.deps) == "string" then
      spec.deps = { spec.deps }
    end
    for _, v in pairs(spec.deps) do
      local p = M.registerpkg(v)
      M.link_dependency(package, p)

      -- if the child package is lazy loaded make sure the child package
      -- is only loaded when the parent package has finished loading
      if package.lazy then
        table.insert(package.on_load, function()
          M.loadtree(package, true)
        end)
        table.insert(p.lazy_load, function(_) end)
      end
    end
  end

  return package
end

--- recurse over all packages and register them
---@param speclist table table of specs
---@param overrides table|nil a package spec that is used to override options
function M.registertree(speclist, overrides)
  overrides = overrides or {}

  -- make sure the overrides override and take into account the packages spec
  overrides = {
    pin = overrides.pin or speclist.pin,
    disable = overrides.disable or speclist.disable
  }

  -- recurse the packages
  for i = 1, #speclist do
    local ok, err = pcall(M.registerpkg, speclist[i], overrides)

    -- if erroring print out the spec and the error
    if not ok then
      error(string.format("%s (spec=%s)", err, vim.inspect(speclist[i])))
    end
  end
end

--- run specified hooks on specified package
---@param package table package spec
---@param type string which hook to run
---@return boolean, string|nil
function M.runhooks(package, type)
  local hooks = package[type]
  if #hooks == 0 then
    return true
  end

  local start = os.clock()

  -- chdir into the package directory to make running external commands
  -- from hooks easier.
  local last_cwd = vim.fn.getcwd()
  vim.fn.chdir(package.dir)

  for i = 1, #hooks do
    local ok, err = pcall(hooks[i])
    if not ok then
      vim.fn.chdir(last_cwd)

      package.error = true
      return false, err
    end
  end

  vim.fn.chdir(last_cwd)
  package.perf[type] = os.clock() - start

  logger:log("hook", "triggered %d %s %s for %s", #hooks, type,
    #hooks == 1 and "hook" or "hooks", package.id)

  return true
end

--- make sure a package has been loaded
---@param package table package
---@param force boolean? force lazy packages to load
---@return boolean|table return true or false if loaded or package spec if lazy loaded
function M.ensureadded(package, force)
  -- print("adding ~ "..package.id)
  local function loadpkg(pkg)
    -- make sure to load the dependencies first
    for _, p in pairs(pkg.dependencies) do
      if not p.loaded then
        M.ensureadded(p, true)
      end
    end

    -- now start loading our plugin
    local start = os.clock()

    -- trigger the packadd for the plugin
    local ok, err = pcall(vim.cmd, "packadd " .. pkg.name)
    if not ok then
      pkg.error = true
      return false, err
    end

    pkg.added = true
    pkg.perf.pack = os.clock() - start
    logger:log("vim", "packadd completed for %s", pkg.id)

    -- set the package to loaded
    package.loaded = true
    logger:log("load", "loaded %s", package.id)

    -- trigger the on_load hooks
    ok, err = M.runhooks(package, "on_load")
    if not ok then
      logger:log("error", "failed to load %s; reason: %s", package.id, err)
      return
    end
  end

  if not package.added and not package.lazy or force then
    loadpkg(package)
  elseif not package.added and package.lazy then
    logger:log("lazy", "registering %d lazy hooks for %s", #package.lazy_load,
      package.id)
    for _, load_cond in pairs(package.lazy_load) do
      -- configure the lazy loader for the user
      local l = require('lazy.utils'):new()
      if l == true then
        logger:log("lazy", "failed to get lazy utils")
        return false
      end
      l:set_load(function()
        logger:log("lazy", "triggered %d lazy hooks for %s", #package.lazy_load,
          package.id)
        loadpkg(package)
      end)

      -- run it
      load_cond(l)
    end
    return package
  end

  return true
end

--- load all packages in package tree
---@param package table package spec table
---@param force boolean? force lazy packages to load
---@return boolean boolean if tree was successfully loaded
function M.loadtree(package, force)
  if not package.exists or not package.enabled or package.error then
    return false
  end

  if package.subtree_loaded then
    return true
  end

  if not package.lazy then
    for i = 1, #package.dependencies do
      if not package.dependencies[i].loaded then
        return false
      end
    end
  end

  if not package.loaded then
    local ok, err = M.ensureadded(package, force)
    if not ok then
      logger:log("error", "failed to load %s; reason: %s", package.id, err)
      return false
    end
  end

  package.subtree_loaded = true

  for i = 1, #package.dependents do
    package.subtree_loaded = M.loadtree(package.dependents[i], force) and package.subtree_loaded
  end

  return package.subtree_loaded
end

--- reload all packages in package table spec
---@param force boolean|nil force all packages to load
function M.reload(force)
  -- cleanup any previous errors in the package table
  for i in pairs(packages) do
    packages[i].error = false
  end

  local reloaded = M.loadtree(root, force)

  if reloaded then
    local ok, err
    M.benchmark("reload", function()
      ok, err = pcall(vim.cmd,
        [[
        silent! helptags ALL
        silent! UpdateRemotePlugins
        ]])
    end)

    if ok then
      logger:log("vim", "reloaded helptags and remote plugins")
    else
      logger:log("error",
        "failed to reload helptags and remote plugins; reason: %s", err)
    end
  end
end

--- check if there's a circular dependency in the package tree
function M.findcycle()
  local index = 0
  local indexes = {}
  local lowlink = {}
  local stack = {}

  -- use tarjan algorithm to find circular dependencies (strongly connected
  -- components)
  local function connect(package)
    indexes[package.id], lowlink[package.id] = index, index
    stack[#stack + 1], stack[package.id] = package, true
    index = index + 1

    for i = 1, #package.dependents do
      local dependent = package.dependents[i]

      if not indexes[dependent.id] then
        local cycle = connect(dependent)
        if cycle then
          return cycle
        else
          lowlink[package.id] = math.min(lowlink[package.id], lowlink[dependent.id])
        end
      elseif stack[dependent.id] then
        lowlink[package.id] = math.min(lowlink[package.id], indexes[dependent.id])
      end
    end

    if lowlink[package.id] == indexes[package.id] then
      local cycle = { package }
      local node

      repeat
        node = stack[#stack]
        stack[#stack], stack[node.id] = nil, nil
        cycle[#cycle + 1] = node
      until node == package

      -- a node is by definition strongly connected to itself ignore single-node
      -- components unless it explicitly specified itself as a dependency
      if #cycle > 2 or package.dependents[package.id] then
        return cycle
      end
    end
  end

  for i = 1, #packages do
    local package = packages[i]

    if not indexes[package.id] then
      local cycle = connect(package)
      if cycle then
        return cycle
      end
    end
  end
end

--- unconfigure a packages tree
---@param package table package to unconfigure
function M.unconfiguretree(package)
  -- unconfigure dependencies
  for i = 1, #package.dependencies do
    package.dependencies[i].subtree_loaded = false
  end

  -- unconfigure dependents
  for i = 1, #package.dependents do
    package.dependents[i].loaded = false
    package.dependents[i].added = false
    package.dependents[i].configured = false

    package.dependents[i].subtree_loaded = false
  end
end

--- update/download a package
---@param package table package spec
---@param cb function callback
function M.sync(package, cb)
  if not package.enabled then
    cb()
    return
  end

  local function log_err(err)
    logger:log("error", "failed to update %s; reason: %s", package.id, err)
  end

  --- configure a package
  ---@param pkg table package spec
  local function configurepkg(pkg)
    M.runhooks(package, "on_config")

    logger:log("config", "package: %s configured", pkg.id)
    package.configured = true
  end

  if package.exists then
    if package.pin then
      cb()
      return
    end

    local function logerr(err)
      logger:log("error", "failed to update %s; reason: %s", package.id, err)
    end

    proc.git_rev_parse(package.dir, "HEAD", function(err, before)
      if err then
        logerr(before)
        cb(err)
      else
        if package.commit then
          proc.git_checkout(package.dir, package.branch, package.commit, function(err, message)
            if err then
              log_err(message)
              cb(err)
            else
              proc.git_rev_parse(package.dir, package.commit, function(err, after)
                if err then
                  log_err(after)
                  cb(err)
                elseif before == after then
                  logger:log("skip", "skipped %s", package.id)
                  cb(err)
                else
                  M.unconfiguretree(package)
                  configurepkg(package)
                  logger:log("update", "updated %s; %s -> %s", package.id, before, after)
                end
              end)
            end
          end)
        else
          proc.git_fetch(package.dir, "origin", package.branch or "HEAD", function(err, message)
            if err then
              log_err(message)
              cb(err)
            else
              proc.git_rev_parse(package.dir, "FETCH_HEAD", function(err, after)
                if err then
                  log_err(after)
                  cb(err)
                elseif before == after then
                  logger:log("skip", "skipped %s", package.id)
                  cb(err)
                else
                  proc.git_reset(package.dir, after, function(err, message)
                    if err then
                      log_err(message)
                    else
                      M.unconfiguretree(package)
                      configurepkg(package)
                      logger:log("update", "updated %s; %s -> %s", package.id, before, after)
                    end

                    cb(err)
                  end)
                end
              end)
            end
          end)
        end
      end
    end)
  else
    logger:log("error", "%s: doesn't exist", package.id)
    proc.git_clone(package.dir, package.url, package.branch, function(err, message)
      if err then
        logger:log("error", "failed to install %s; reason: %s",
          package.id, message)
      else
        if package.commit then
          proc.git_checkout(package.dir, package.branch, package.commit, function(err, message)
            if err then
              logger:log("error", "failed to checkout %s; reason: %s", package.id, message)
            else
              package.exists = true
              M.unconfiguretree(package)
              configurepkg(package)
              logger:log("install", "installed %s", package.id)
            end
          end)
        else
          package.exists = true
          M.unconfiguretree(package)
          configurepkg(package)
          logger:log("install", "installed %s", package.id)
        end
      end

      cb(err)
    end)
  end
end

--- sync a tree of plugins
---@param tree table tree of plugins
---@param cb function? callback
function M.synctree(tree, cb)
  local progress = 0
  local has_errors = false

  local function done(err)
    progress = progress + 1
    has_errors = has_errors or err

    if progress == #tree then
      -- TODO: implement clean
      -- clean()
      M.reload()

      if has_errors then
        logger:log("error", "there were errors during sync; see :messages or :DepLog for more information")
      else
        logger:log("update", "synchronized %s %s", #tree, #tree == 1 and "package" or "packages")
      end

      if cb then
        cb()
      end
    end
  end

  for i in pairs(tree) do
    M.sync(tree[i], done)
  end
end


local function print_list(cb)
  local function get_commits(cb)
    local results = {}
    local done = 0
    for i = 1, #packages do
      local package = packages[i]

      if package.exists then
        proc.git_rev_parse(package.dir, "HEAD", function(err, commit)
          if not err then
            results[package.id] = commit
          end

          done = done + 1
          if done == #packages then
            cb(results)
          end
        end)
      else
        done = done + 1
      end
    end
  end

  get_commits(function(commits)
    local buffer = vim.api.nvim_create_buf(true, true)
    local line, indent = 0, 0

    local function print(chunks)
      local concat = {}
      local column = 0

      for _ = 1, indent do
        concat[#concat + 1] = "  "
        column = column + 2
      end

      if not chunks then
        chunks = {}
      elseif type(chunks) == "string" then
        chunks = { { chunks } }
      end

      for i = 1, #chunks do
        local chunk = chunks[i]
        concat[#concat + 1] = chunk[1]
        chunk.offset, column = column, column + #chunk[1]
      end

      vim.api.nvim_buf_set_lines(buffer, line, -1, false, { table.concat(concat) })

      for i = 1, #chunks do
        local chunk = chunks[i]
        if chunk[2] then
          vim.api.nvim_buf_add_highlight(buffer, -1, chunk[2], line, chunk.offset, chunk.offset + #chunk[1])
        end
      end

      line = line + 1
    end

    print(string.format("Installed packages (%s):", #packages))
    indent = 1

    local loaded = {}

    local function dry_load(package)
      if loaded[package.id] then
        return
      end

      for i = 1, #package.dependencies do
        if not loaded[package.dependencies[i].id] then
          return
        end
      end

      loaded[package.id], loaded[#loaded + 1] = true, package

      local chunk = {
        { string.format("[%s] ", commits[package.id] or "       "), "Comment" },
        { package.id, "Underlined" },
      }

      if not package.exists then
        chunk[#chunk + 1] = { " *not installed", "Comment" }
      end

      if not package.loaded then
        chunk[#chunk + 1] = { " *not loaded", "Comment" }
      end

      if not package.enabled then
        chunk[#chunk + 1] = { " *disabled", "Comment" }
      end

      if package.pin then
        chunk[#chunk + 1] = { " *pinned", "Comment" }
      end

      print(chunk)

      for i = 1, #package.dependents do
        dry_load(package.dependents[i])
      end
    end

    dry_load(root)
    indent = 0

    print()
    print("Load time (μs):")
    indent = 1
    local profiles = {}

    for i = 1, #packages do
      local package = packages[i]
      local profile = {
        package = package,
        total = 0,
        setup = package.perf.on_setup or 0,
        load = package.perf.on_load or 0,
        pack = package.perf.pack or 0,

        "total",
        "setup",
        "pack",
        "load",
      }

      if package == root then
        for k, v in pairs(perf) do
          if profile[k] then
            profile[k] = profile[k] + v
          end
        end
      end

      for j = 1, #profile do
        profile.total = profile.total + profile[profile[j]]
      end

      profiles[#profiles + 1] = profile
    end

    table.sort(profiles, function(a, b)
      return a.total > b.total
    end)

    for i = 1, #profiles do
      local profile = profiles[i]
      local chunk = {
        { "- ", "Comment" },
        { profile.package.id, "Underlined" },
        { string.rep(" ", 40 - #profile.package.id) },
      }

      for j = 1, #profile do
        local key, value = profile[j], profile[profile[j]]
        chunk[#chunk + 1] = { string.format(" %5s ", key), "Comment" }
        chunk[#chunk + 1] = { string.format("%4d", value * 1000000) }
      end

      print(chunk)
    end

    indent = 0
    print()
    print("Dependency graph:")

    local function walk_graph(package)
      local chunk = {
        { "| ", "Comment" },
        { package.id, "Underlined" },
      }

      local function add_edges(p)
        for i = 1, #p.dependencies do
          local dependency = p.dependencies[i]

          if dependency ~= root and not chunk[dependency.id] then -- don't convolute the list
            chunk[#chunk + 1] = { " " .. dependency.id, "Comment" }
            chunk[dependency.id] = true
            add_edges(dependency)
          end
        end
      end

      add_edges(package)
      print(chunk)

      for i = 1, #package.dependents do
        indent = indent + 1
        walk_graph(package.dependents[i])
        indent = indent - 1
      end
    end

    walk_graph(root)

    -- print()
    -- print("Debug information:")

    -- local debug = {}
    -- for l in vim.inspect(packages):gmatch("[^\n]+") do
    --   debug[#debug + 1] = l
    -- end

    -- vim.api.nvim_buf_set_lines(buffer, line, -1, false, debug)
    -- vim.api.nvim_buf_set_name(buffer, "packages.dep")
    -- vim.api.nvim_buf_set_option(buffer, "bufhidden", "wipe")
    -- vim.api.nvim_buf_set_option(buffer, "modifiable", false)

    vim.cmd("vsp")
    vim.api.nvim_win_set_buf(0, buffer)

    if cb then
      cb()
    end
  end)
end

-- basically the main function of our program
return function(opts)
  logger.pipe = logger:setup()

  --- make comparison for table.sort
  ---@param a table package spec a
  ---@param b table package spec b
  ---@return boolean
  local function comp(a, b)
    -- NOTE: this doesn't have to be in any real order, it just has to be
    -- consistant, thus we can just check if the unicode value of one package
    -- id is less than the other
    return a.id < b.id
  end

  initialized, err = pcall(function()
    base_dir = opts.base_dir or vim.fn.stdpath("data").."/site/pack/deps/opt/"
    M.benchmark("load", function()
      -- register all packages
      root = M.registerpkg("squibid/dep")
      M.registertree(opts)

      -- sort packages
      table.sort(packages, comp)

      -- sort package dependencies
      for i = 1, #packages do
        table.sort(packages[i].dependencies, comp)
        table.sort(packages[i].dependents, comp)
      end

      -- make sure there arent any circular dependencies
      M.findcycle()
    end)

    -- load packages
    M.reload()

    --- check if a package should be synced
    ---@param package table package table spec
    ---@return boolean sync
    local function shouldsync(package)
      if opts.sync == "new" or opts.sync == nil then
        return not package.exists
      else
        return opts.sync == "always"
      end
    end

    -- get all package that need syncing
    local targets = {}
    for i in pairs(packages) do
      if shouldsync(packages[i]) then
        targets[i] = packages[i]
      end
    end

    M.synctree(targets)
  end)

  if not initialized then
    logger:log("error", err)
  end

  vim.api.nvim_create_user_command("DepLog", function()
    vim.cmd('vsp '..logger.path)
    vim.opt_local.readonly = true

    local w = vim.uv.new_fs_event()
    local function watch_file(fname)
      local fullpath = vim.api.nvim_call_function(
        'fnamemodify', { fname, ':p' })
      w:start(fullpath, {}, vim.schedule_wrap(function(...)
        vim.api.nvim_command('checktime')
        -- Debounce: stop/start.
        w:stop()
        watch_file(fname)
      end))
    end

    watch_file(logger.path)
  end, {})

  vim.api.nvim_create_user_command("DepSync", function()
    M.synctree(packages)
  end, {})

  vim.api.nvim_create_user_command("DepReload", function()
    M.reload()
  end, {})

  vim.api.nvim_create_user_command("DepList", function()
    print_list()
  end, {})

  logger:cleanup()
end