nvim/lua/git.lua

local util = require("util")

local HIGHLIGHTS = {
    GitDeleted = "Statement",
    GitDirty = "Statement",
    GitIgnored = "Comment",
    GitMerge = "Constant",
    GitNew = "PreProc",
    GitRenamed = "PreProc",
    GitStaged = "Constant",
}

local UNMERGED = {
    DD = true,
    AU = true,
    UD = true,
    UA = true,
    DU = true,
    AA = true,
    UU = true,
}

---@param code string
---@return string?
local function format(code)
    if code == "" then
        return nil
    end
    local char, hl
    if code == "??" then
        char, hl = "?", "GitNew"
    elseif code == "!!" then
        char, hl = "!", "GitIgnored"
    elseif UNMERGED[code] then
        char, hl = "U", "GitMerge"
    else
        local x, y = code:sub(1, 1), code:sub(2, 2)
        if x == "R" or y == "R" then
            char, hl = "R", "GitRenamed"
        elseif y == "M" or y == "T" then
            char, hl = "M", "GitDirty"
        elseif y == "D" then
            char, hl = "D", "GitDeleted"
        elseif y == " " and x == "D" then
            char, hl = "D", "GitStaged"
        elseif y == " " and x == "A" then
            char, hl = "A", "GitStaged"
        elseif y == " " and x ~= " " then
            char, hl = "M", "GitStaged"
        else
            char, hl = "M", "GitDirty"
        end
    end
    return string.format("%%#%s#%s%%*", hl, char)
end

---@param path string
---@return string? gitdir
---@return string? worktree
local function resolve(path)
    local found = vim.fs.find(".git", { upward = true, path = path })[1]
    if not found then
        return nil
    end
    local worktree = vim.fs.dirname(found)
    local stat = vim.uv.fs_stat(found)
    if not stat then
        return nil
    end
    if stat.type == "directory" then
        return found, worktree
    end
    local f = io.open(found, "r")
    if not f then
        return nil
    end
    local content = f:read("*a")
    f:close()
    local gitdir = content:match("gitdir:%s*(%S+)")
    if not gitdir then
        return nil
    end
    if not gitdir:match("^/") then
        gitdir = vim.fs.joinpath(worktree, gitdir)
    end
    return vim.fs.normalize(gitdir), worktree
end

---@class ow.Git.Repo
---@field gitdir string
---@field worktree string
---@field buffers integer[]
---@field watcher? uv.uv_fs_event_t
---@field refresh ow.Util.Debouncer
local Repo = {}
Repo.__index = Repo

function Repo:start_watcher()
    local watcher = assert(vim.uv.new_fs_event())
    assert(watcher:start(self.gitdir, {}, function(err, filename)
        if err or (filename ~= "index" and filename ~= "HEAD") then
            return
        end
        self.refresh()
    end))
    self.watcher = watcher
end

function Repo:stop_watcher()
    self.refresh:cancel()
    if self.watcher then
        self.watcher:stop()
        self.watcher:close()
        self.watcher = nil
    end
end

---@param buf integer
function Repo:add_buffer(buf)
    table.insert(self.buffers, buf)
end

---@param buf integer
function Repo:remove_buffer(buf)
    for i, b in ipairs(self.buffers) do
        if b == buf then
            table.remove(self.buffers, i)
            break
        end
    end
end

---@param repo ow.Git.Repo
local function do_refresh(repo)
    vim.system({
        "git",
        "-c",
        "core.quotePath=false",
        "status",
        "--porcelain=v1",
    }, { cwd = repo.worktree, text = true }, function(obj)
        vim.schedule(function()
            local statuses = {}
            if obj.code == 0 then
                for line in (obj.stdout or ""):gmatch("[^\r\n]+") do
                    local code = line:sub(1, 2)
                    local path_part = line:sub(4)
                    local arrow = path_part:find(" -> ", 1, true)
                    if arrow then
                        path_part = path_part:sub(arrow + 4)
                    end
                    statuses[vim.fs.joinpath(repo.worktree, path_part)] =
                        format(code)
                end
            end
            for _, buf in ipairs(repo.buffers) do
                if vim.api.nvim_buf_is_valid(buf) then
                    local status = statuses[vim.api.nvim_buf_get_name(buf)]
                    if vim.b[buf].git_status ~= status then
                        vim.b[buf].git_status = status
                        vim.cmd.redrawstatus({ bang = true })
                    end
                end
            end
            vim.api.nvim_exec_autocmds("User", {
                pattern = "GitRefresh",
                data = { gitdir = repo.gitdir, worktree = repo.worktree },
            })
        end)
    end)
end

---@param gitdir string
---@param worktree string
---@return ow.Git.Repo
function Repo.new(gitdir, worktree)
    local self = setmetatable({
        gitdir = gitdir,
        worktree = worktree,
        buffers = {},
    }, Repo)
    self.refresh = util.debounce(function()
        do_refresh(self)
    end, 50)
    self:start_watcher()
    return self
end

---@type table<string, ow.Git.Repo>
local repo_by_gitdir = {}

---@type table<integer, ow.Git.Repo>
local repo_by_buf = {}

---@param buf integer
---@return ow.Git.Repo?
local function register(buf)
    local existing = repo_by_buf[buf]
    if existing then
        return existing
    end
    local path = vim.api.nvim_buf_get_name(buf)
    if path == "" then
        return nil
    end
    local gitdir, worktree = resolve(path)
    if not gitdir or not worktree then
        return nil
    end
    local repo = repo_by_gitdir[gitdir]
    if not repo then
        repo = Repo.new(gitdir, worktree)
        repo_by_gitdir[gitdir] = repo
    end
    repo:add_buffer(buf)
    repo_by_buf[buf] = repo
    return repo
end

---@param buf integer
local function unregister(buf)
    local repo = repo_by_buf[buf]
    if not repo then
        return
    end
    repo_by_buf[buf] = nil
    repo:remove_buffer(buf)
    if #repo.buffers == 0 then
        repo:stop_watcher()
        repo_by_gitdir[repo.gitdir] = nil
    end
end

---@param buf integer
local function refresh(buf)
    if not vim.api.nvim_buf_is_valid(buf) or vim.bo[buf].buftype ~= "" then
        return
    end
    local repo = register(buf)
    if not repo then
        vim.b[buf].git_status = nil
        return
    end
    repo.refresh()
end

local M = {}

function M.status()
    return vim.b.git_status or ""
end

function M.setup()
    for name, link in pairs(HIGHLIGHTS) do
        vim.api.nvim_set_hl(0, name, { link = link, default = true })
    end
    local group = vim.api.nvim_create_augroup("ow.git", { clear = true })
    vim.api.nvim_create_autocmd(
        { "BufReadPost", "BufNewFile", "BufWritePost", "FileChangedShellPost" },
        {
            group = group,
            callback = function(args)
                refresh(args.buf)
            end,
        }
    )
    vim.api.nvim_create_autocmd({ "BufDelete", "BufWipeout" }, {
        group = group,
        callback = function(args)
            unregister(args.buf)
        end,
    })
    vim.api.nvim_create_autocmd("FocusGained", {
        group = group,
        callback = function()
            refresh(vim.api.nvim_get_current_buf())
        end,
    })
    vim.api.nvim_create_autocmd("VimLeavePre", {
        group = group,
        callback = function()
            for _, repo in pairs(repo_by_gitdir) do
                repo:stop_watcher()
            end
        end,
    })
end

return M