nvim.config/lua/scripts/go-gen-constructor.lua

local M = {}

-- Get the struct name from the current line
local function get_struct_name()
    local line = vim.api.nvim_get_current_line()

    -- Try to match: type StructName struct
    local struct_name = line:match("^type%s+([%w_]+)%s+struct")

    -- If not found, try to match if cursor is on struct definition
    if not struct_name then
        struct_name = line:match("type%s+([%w_]+)%s+struct")
    end

    return struct_name
end

-- Check if constructor already exists and find its location
local function find_existing_constructor(constructor_name, bufnr)
    local lines = vim.api.nvim_buf_get_lines(bufnr, 0, -1, false)

    for i, line in ipairs(lines) do
        -- Look for function definition with constructor name
        if line:match("^func%s+" .. constructor_name .. "%s*%(") then
            local start_line = i
            local end_line = i

            -- Find the end of the function (matching braces)
            local bracket_count = 0
            for j = i, #lines do
                for char in lines[j]:gmatch(".") do
                    if char == "{" then
                        bracket_count = bracket_count + 1
                    elseif char == "}" then
                        bracket_count = bracket_count - 1
                    end
                end

                if bracket_count == 0 and lines[j]:match("}") then
                    end_line = j
                    break
                end
            end
            return { start_line = start_line, end_line = end_line, line = line }
        end
    end

    return nil
end

-- Extract return type from existing constructor
local function extract_return_type_from_existing(constructor_line)
    -- Match: func ConstructorName(params) ReturnType {
    local return_type = constructor_line:match("%)%s*([%w%.]+)%s*{")
    return return_type
end

local function remove_existing_constructor(constructor_name, bufnr)
    local location = find_existing_constructor(constructor_name, bufnr)
    if location then
        -- Remove the constructor
        vim.api.nvim_buf_set_lines(bufnr, location.start_line - 1, location.end_line, false, {})
        return location
    end
    return nil
end

-- Find interface that the struct should implement
local function find_interface_for_struct(struct_name, bufnr)
    local lines = vim.api.nvim_buf_get_lines(bufnr, 0, -1, false)

    -- Look for interfaces defined in the file
    for i, line in ipairs(lines) do
        -- Match interface definition
        local interface_name = line:match("^type%s+([%w_]+)%s+interface")
        if interface_name then
            -- Check if this interface has methods that match our struct
            -- Simple check: if interface name contains struct name or vice versa
            if
                string.find(interface_name, struct_name)
                or string.find(struct_name, interface_name)
                or string.lower(struct_name) == string.lower(interface_name)
            then
                return interface_name
            end
        end
    end

    -- If no matching interface found, check if struct has an interface with same name + "er" or "Service"
    local possible_interfaces = {
        string.gsub(struct_name, "Service$", "Service"),
        struct_name .. "er",
        "I" .. struct_name,
        string.gsub(struct_name, "^%l", string.upper),
    }

    for _, interface_name in ipairs(possible_interfaces) do
        for i, line in ipairs(lines) do
            if line:match("^type%s+" .. interface_name .. "%s+interface") then
                return interface_name
            end
        end
    end

    return nil
end

-- Parse struct fields
local function parse_struct_fields(struct_name, bufnr)
    local fields = {}
    local lines = vim.api.nvim_buf_get_lines(bufnr, 0, -1, false)
    local in_struct = false
    local bracket_count = 0
    local start_line = 0

    for i, line in ipairs(lines) do
        -- Find the struct definition
        if line:match("^type%s+" .. struct_name .. "%s+struct") then
            in_struct = true
            start_line = i
            -- Check if opening brace is on same line
            if line:match("{") then
                bracket_count = 1
            end
            if line:match("}") then
                break
            end
        elseif in_struct then
            -- Count brackets to know when we're done
            if not line:match("^%s*//") then -- Skip comment lines for bracket counting
                for char in line:gmatch(".") do
                    if char == "{" then
                        bracket_count = bracket_count + 1
                    elseif char == "}" then
                        bracket_count = bracket_count - 1
                    end
                end
            end

            -- Parse field if we're inside the struct and it's not the closing brace
            if bracket_count > 0 and i > start_line then
                -- Remove comments and trim
                local clean_line = line:gsub("//.*$", ""):gsub("^%s*(.-)%s*$", "%1")

                -- Skip empty lines and closing braces
                if clean_line ~= "" and not clean_line:match("^}") and not clean_line:match("^{") then
                    -- Extract field name and type
                    -- Handle multiple fields with same type: field1, field2 Type
                    local field_part = clean_line

                    -- Check if line contains a type declaration
                    if field_part:match("[%w_]+%s+[%w%.%*%[%]]+$") then
                        -- Split by last space to get type and field names
                        local last_space = field_part:match("^.*()%s[%w%.%*%[%]]+$")
                        if last_space then
                            local type_part = field_part:sub(last_space + 1)
                            local fields_part = field_part:sub(1, last_space - 1)

                            -- Split multiple field names by comma
                            for field_name in fields_part:gmatch("([%w_]+),?%s*") do
                                if field_name ~= "" then
                                    table.insert(fields, {
                                        name = field_name,
                                        type = type_part,
                                        line = clean_line,
                                    })
                                end
                            end
                        end
                    end
                end
            end

            -- Stop parsing when we exit the struct
            if bracket_count == 0 then
                break
            end
        end
    end

    return fields
end

-- Generate constructor function
local function generate_constructor(struct_name, interface_name, fields, existing_constructor)
    local constructor_name = "New" .. struct_name
    if string.match(struct_name, "^[a-z]") and interface_name then
        constructor_name = "New" .. interface_name
    else
        constructor_name = "New" .. string.upper(string.sub(struct_name, 1, 1)) .. string.sub(struct_name, 2)
    end

    local lines = {}

    -- Function signature
    local return_type
    if existing_constructor and existing_constructor.line then
        return_type = extract_return_type_from_existing(existing_constructor.line)
    end

    if not return_type then
        return_type = interface_name or struct_name
    end

    -- Build parameters
    local params = {}
    local field_params = {}
    for _, field in ipairs(fields) do
        local param_name = field.name
        local param_type = field.type

        table.insert(params, param_name .. " " .. param_type)
        field_params[param_name] = param_name
    end

    -- Build function signature
    local func_sig = "func " .. constructor_name .. "("
    if #params > 0 then
        func_sig = func_sig .. table.concat(params, ", ")
    end
    func_sig = func_sig .. ") " .. return_type .. " {"
    table.insert(lines, func_sig)

    -- Function body
    table.insert(lines, "\treturn &" .. struct_name .. "{")

    -- Field assignments
    local field_assignments = {}
    for _, field in ipairs(fields) do
        local value = field_params[field.name] or field.name
        table.insert(field_assignments, "\t\t" .. field.name .. ": " .. value .. ",")
    end

    for _, assignment in ipairs(field_assignments) do
        table.insert(lines, assignment)
    end

    table.insert(lines, "\t}")
    table.insert(lines, "}")

    return lines, constructor_name
end

-- Main function to generate constructor
function M.generate_constructor()
    local bufnr = vim.api.nvim_get_current_buf()

    -- Check if we're in a Go file
    local filetype = vim.bo[bufnr].filetype
    if filetype ~= "go" then
        vim.notify("This command only works in Go files", vim.log.levels.ERROR)
        return
    end

    -- Get struct name under cursor
    local struct_name = get_struct_name()
    if not struct_name then
        vim.notify("Cursor is not on a struct definition", vim.log.levels.ERROR)
        return
    end

    -- Find interface if exists
    local interface_name = find_interface_for_struct(struct_name, bufnr)

    -- Parse struct fields
    local fields = parse_struct_fields(struct_name, bufnr)

    if #fields == 0 then
        vim.notify("No fields found in struct " .. struct_name, vim.log.levels.WARN)
        return
    end

    -- Check for existing constructor and remove it (but save its info)
    local constructor_name = "New" .. struct_name
    if string.match(struct_name, "^[a-z]") and interface_name then
        constructor_name = "New" .. interface_name
    else
        constructor_name = "New" .. string.upper(string.sub(struct_name, 1, 1)) .. string.sub(struct_name, 2)
    end

    -- Remove existing constructor if it exists
    local removed = remove_existing_constructor(constructor_name, bufnr)

    -- Generate constructor
    local constructor_lines = generate_constructor(struct_name, interface_name, fields, removed)

    -- Find where to insert the constructor (after the struct)
    local lines = vim.api.nvim_buf_get_lines(bufnr, 0, -1, false)
    local insert_line = 0

    for i, line in ipairs(lines) do
        if line:match("^type%s+" .. struct_name .. "%s+struct") then
            -- Find the end of the struct
            local bracket_count = 0
            for j = i, #lines do
                local current_line = lines[j]
                -- Skip comment lines for bracket counting
                if not current_line:match("^%s*//") then
                    for char in current_line:gmatch(".") do
                        if char == "{" then
                            bracket_count = bracket_count + 1
                        elseif char == "}" then
                            bracket_count = bracket_count - 1
                        end
                    end
                end

                if bracket_count == 0 then
                    insert_line = j
                    break
                end
            end
            break
        end
    end

    if insert_line > 0 then
        -- Check if we need to add an empty line
        if insert_line < #lines and lines[insert_line + 1] and lines[insert_line + 1]:match("^%s*$") then
            -- Next line is already empty, insert after it
            vim.api.nvim_buf_set_lines(bufnr, insert_line + 1, insert_line + 1, false, constructor_lines)
        else
            -- Add an empty line before constructor
            table.insert(constructor_lines, 1, "")
            vim.api.nvim_buf_set_lines(bufnr, insert_line, insert_line, false, constructor_lines)
        end

        if removed then
            vim.notify("Replaced existing constructor: " .. constructor_name, vim.log.levels.INFO)
        else
            vim.notify("Generated constructor: " .. constructor_name, vim.log.levels.INFO)
        end
    else
        vim.notify("Could not find where to insert constructor", vim.log.levels.ERROR)
    end
end

-- Setup function
function M.setup()
    -- Create command
    vim.api.nvim_create_user_command("GoGenerateConstructor", M.generate_constructor, {
        desc = "Generate constructor for Go struct under cursor",
    })

    -- Optional keybinding
    vim.keymap.set("n", "<leader>cec", M.generate_constructor, {
        desc = "Generate constructor for struct",
        silent = true,
    })
end

return M