nvim/lua/ts.lua

local log = require("log")

local M = {}

---@class ow.TS.ParserSpec
---@field lang string
---@field location? string
---@field generate? boolean
---@field from_json? boolean

---@class ow.TS.SingleParser
---@field [1] string
---@field name? string
---@field version? string | vim.VersionRange
---@field lang? string inferred from URL (strip "tree-sitter-" prefix) if omitted
---@field location? string
---@field generate? boolean
---@field from_json? boolean

---@class ow.TS.MultiParser
---@field [1] string
---@field name? string
---@field version? string | vim.VersionRange
---@field specs ow.TS.ParserSpec[]

---@alias ow.TS.ParserEntry string | ow.TS.SingleParser | ow.TS.MultiParser

---@class ow.TS.Parser
---@field name string
---@field path string
---@field specs ow.TS.ParserSpec[]

---@param path string
---@param lang string
---@return string
local function parser_so(path, lang)
    return vim.fs.joinpath(path, "parser", lang .. ".so")
end

---@param buf integer
local function start_treesitter(buf)
    local ok, err = pcall(vim.treesitter.start, buf)
    if not ok then
        log.error("Failed to enable treesitter for buffer %d: %s", buf, err)
        return
    end
    for _, win in ipairs(vim.fn.win_findbuf(buf)) do
        vim.wo[win].foldmethod = "expr"
        vim.wo[win].foldexpr = "v:lua.vim.treesitter.foldexpr()"
    end
end

---@param lang string
local function activate_open_buffers(lang)
    local fts = vim.treesitter.language.get_filetypes(lang)
    for _, buf in ipairs(vim.api.nvim_list_bufs()) do
        if
            vim.api.nvim_buf_is_loaded(buf)
            and vim.list_contains(fts, vim.bo[buf].filetype)
        then
            start_treesitter(buf)
        end
    end
end

---@param parser ow.TS.Parser
---@param spec ow.TS.ParserSpec
function M.build(parser, spec)
    local cwd = spec.location and vim.fs.joinpath(parser.path, spec.location)
        or parser.path
    local out = parser_so(parser.path, spec.lang)
    vim.fn.mkdir(vim.fs.dirname(out), "p")

    local function on_build(r)
        if r.code ~= 0 then
            log.error(
                "Failed to build parser for %s: %s",
                spec.lang,
                r.stderr or ""
            )
            return
        end
        vim.treesitter.language.add(spec.lang, { path = out })
        activate_open_buffers(spec.lang)
    end

    local function do_build()
        vim.system(
            { "tree-sitter", "build", "-o", out },
            { cwd = cwd },
            vim.schedule_wrap(on_build)
        )
    end

    if spec.generate then
        local cmd = {
            "tree-sitter",
            "generate",
            "--abi",
            tostring(vim.treesitter.language_version),
        }
        if spec.from_json ~= false then
            table.insert(cmd, "src/grammar.json")
        end
        vim.system(
            cmd,
            {
                cwd = cwd,
                env = { TREE_SITTER_JS_RUNTIME = "native" },
            },
            vim.schedule_wrap(function(r)
                if r.code ~= 0 then
                    log.error(
                        "Failed to generate parser for %s: %s",
                        spec.lang,
                        r.stderr or ""
                    )
                    return
                end
                do_build()
            end)
        )
    else
        do_build()
    end
end

---@param pack_name string
---@return string
local function lang_from_name(pack_name)
    local lang = pack_name:match("^tree%-sitter%-(.+)$")
    if not lang then
        error("cannot derive lang from pack name: " .. pack_name)
    end
    return lang
end

---@param entry ow.TS.ParserEntry
---@return vim.pack.Spec
local function entry_pack_spec(entry)
    if type(entry) == "string" then
        return { src = entry }
    end
    return {
        src = entry[1],
        name = entry.name,
        version = entry.version,
    }
end

---@param entry ow.TS.ParserEntry
---@param pack_name string fallback for lang when the entry doesn't set it
---@return ow.TS.ParserSpec[]
local function entry_specs(entry, pack_name)
    if type(entry) == "string" then
        return { { lang = lang_from_name(pack_name) } }
    end
    if entry.specs then
        return entry.specs
    end
    ---@cast entry ow.TS.SingleParser
    return {
        {
            lang = entry.lang or lang_from_name(pack_name),
            location = entry.location,
            generate = entry.generate,
            from_json = entry.from_json,
        },
    }
end

---@param languages string[]
---@return string[]
local function collect_filetypes(languages)
    local seen = {}
    for _, lang in ipairs(languages) do
        for _, ft in ipairs(vim.treesitter.language.get_filetypes(lang)) do
            seen[ft] = true
        end
    end
    return vim.tbl_keys(seen)
end

---@param opts { parsers: ow.TS.ParserEntry[], languages: string[] }
function M.setup(opts)
    ---@type vim.pack.Spec[]
    local pack_specs = {}
    ---@type table<string, ow.TS.ParserEntry>
    local entries_by_src = {}
    for _, entry in ipairs(opts.parsers) do
        local pack_spec = entry_pack_spec(entry)
        table.insert(pack_specs, pack_spec)
        entries_by_src[pack_spec.src] = entry
    end

    ---@type ow.TS.Parser[]
    local parsers = {}
    local changed = require("pack").install(pack_specs, function(data)
        if not data.spec.name then
            log.error("Missing name for parser: %s", data.spec.src)
            return
        end
        local entry = entries_by_src[data.spec.src]
        if not entry then
            return
        end
        table.insert(parsers, {
            name = data.spec.name,
            path = data.path,
            specs = entry_specs(entry, data.spec.name),
        })
    end)

    for _, parser in ipairs(parsers) do
        for _, spec in ipairs(parser.specs) do
            local so = parser_so(parser.path, spec.lang)
            if changed[parser.name] or not vim.uv.fs_stat(so) then
                M.build(parser, spec)
            else
                vim.treesitter.language.add(spec.lang, { path = so })
            end
        end
    end

    vim.api.nvim_create_autocmd("FileType", {
        pattern = collect_filetypes(opts.languages),
        callback = function(ev)
            start_treesitter(ev.buf)
        end,
    })
end

return M