nvim/lua/git/status.lua

local M = {}

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

---@alias ow.Git.Status.EntryKind "untracked"|"unstaged"|"staged"|"unmerged"|"ignored"

---@class ow.Git.Status.Entry
---@field path string
---@field kind ow.Git.Status.EntryKind
---@field char string
---@field hl string
---@field orig string?

---@class ow.Git.Status.BranchInfo
---@field head string?
---@field upstream string?
---@field ahead integer
---@field behind integer

---@class ow.Git.Status
---@field branch ow.Git.Status.BranchInfo
---@field entries table<string, ow.Git.Status.Entry[]>
local Status = {}
Status.__index = Status

---@param kind ow.Git.Status.EntryKind
---@return ow.Git.Status.Entry[]
function Status:by_kind(kind)
    local out = {}
    for _, list in pairs(self.entries) do
        for _, e in ipairs(list) do
            if e.kind == kind then
                table.insert(out, e)
            end
        end
    end
    return out
end

---@param prefix string
---@return ow.Git.Status.Entry[]
function Status:aggregate_at(prefix)
    local match = (prefix == "" or prefix == ".") and "" or prefix .. "/"
    local seen = {}
    local out = {}
    for path, list in pairs(self.entries) do
        if path == prefix or vim.startswith(path, match) then
            for _, e in ipairs(list) do
                local key = e.char .. "\0" .. e.hl
                if not seen[key] then
                    seen[key] = true
                    table.insert(out, e)
                end
            end
        end
    end
    table.sort(out, function(a, b)
        return a.char < b.char
    end)
    return out
end

---@param x string
---@return string char, string hl
local function staged_attrs(x)
    if x == "R" then
        return "R", "GitRenamed"
    end
    return x, "GitStaged"
end

---@param y string
---@return string char, string hl
local function unstaged_attrs(y)
    if y == "R" then
        return "R", "GitRenamed"
    end
    if y == "D" then
        return "D", "GitDeleted"
    end
    return y, "GitUnstaged"
end

---@param line string
---@return ow.Git.Status.BranchInfo
local function parse_branch_line(line)
    local info = { ahead = 0, behind = 0 }
    local content = line:sub(4)
    local arrow = content:find("...", 1, true)
    if not arrow then
        info.head = content
        return info
    end
    info.head = content:sub(1, arrow - 1)
    local rest = content:sub(arrow + 3)
    local bracket = rest:find(" %[")
    if not bracket then
        info.upstream = rest
        return info
    end
    info.upstream = rest:sub(1, bracket - 1)
    local inside = rest:match("%[([^%]]+)%]")
    if inside then
        info.ahead = (tonumber(inside:match("ahead (%d+)")) or 0) --[[@as integer]]
        info.behind = (tonumber(inside:match("behind (%d+)")) or 0) --[[@as integer]]
    end
    return info
end

---@param line string
---@return string x, string y, string path, string? orig
local function parse_status_line(line)
    local x = line:sub(1, 1)
    local y = line:sub(2, 2)
    local rest = line:sub(4)
    local orig
    if x == "R" or x == "C" or y == "R" or y == "C" then
        local arrow = rest:find(" -> ", 1, true)
        if arrow then
            orig = rest:sub(1, arrow - 1)
            rest = rest:sub(arrow + 4)
        end
    end
    return x, y, rest, orig
end

---@param entries table<string, ow.Git.Status.Entry[]>
---@param entry ow.Git.Status.Entry
local function add(entries, entry)
    local key = entry.path
    if key:sub(-1) == "/" then
        key = key:sub(1, -2)
    end
    local list = entries[key] or {}
    table.insert(list, entry)
    entries[key] = list
end

---@param stdout string
---@return ow.Git.Status
function M.parse(stdout)
    local branch = { ahead = 0, behind = 0 }
    ---@type table<string, ow.Git.Status.Entry[]>
    local entries = {}
    for line in stdout:gmatch("[^\r\n]+") do
        if line:sub(1, 2) == "##" then
            branch = parse_branch_line(line)
        else
            local x, y, path, orig = parse_status_line(line)
            if x == "?" and y == "?" then
                add(entries, {
                    path = path,
                    kind = "untracked",
                    char = "?",
                    hl = "GitUntracked",
                })
            elseif x == "!" and y == "!" then
                add(entries, {
                    path = path,
                    kind = "ignored",
                    char = "i",
                    hl = "GitIgnored",
                })
            elseif UNMERGED[x .. y] then
                add(entries, {
                    path = path,
                    kind = "unmerged",
                    char = "!",
                    hl = "GitUnmerged",
                })
            else
                if x ~= " " then
                    local char, hl = staged_attrs(x)
                    add(entries, {
                        path = path,
                        kind = "staged",
                        char = char,
                        hl = hl,
                        orig = orig,
                    })
                end
                if y ~= " " then
                    local char, hl = unstaged_attrs(y)
                    add(entries, {
                        path = path,
                        kind = "unstaged",
                        char = char,
                        hl = hl,
                        orig = orig,
                    })
                end
            end
        end
    end
    return setmetatable({ branch = branch, entries = entries }, Status)
end

return M