From eb35a87cdb324fbbaf433ed008d5a7fa43056a8d Mon Sep 17 00:00:00 2001
From: Henry <henrysg88@gmail.com>
Date: Fri, 19 Apr 2024 23:18:38 -0400
Subject: [PATCH] Update GenConstructor function

---
 lua/gopher/gen_constructor.lua | 94 +++++++++++++++++++---------------
 1 file changed, 52 insertions(+), 42 deletions(-)

diff --git a/lua/gopher/gen_constructor.lua b/lua/gopher/gen_constructor.lua
index 81ef477..80cb1f5 100644
--- a/lua/gopher/gen_constructor.lua
+++ b/lua/gopher/gen_constructor.lua
@@ -1,51 +1,61 @@
 local u = require "gopher._utils"
+local ts_utils = require "gopher._utils.ts.init"
 
-return function (cmd_args)
-    local parser = vim.treesitter.get_parser()
-    if not parser then
-        u.notify("no parser not available", "error")
-        return
-    end
+return function(_)
+  -- TODO: allow to pass a funcName as an argument
 
-    -- When the command is runned on the 'struct name' the current node will be
-    -- the node referring directly to the 'struct name', therefore, to get the
-    -- node that refers to the entire struct we have to get the parent.
-    local current_node = vim.treesitter.get_node():parent()
-    if not current_node then
-        u.notify("no node found at cursor", "error")
-        return
-    end
+  local row, col = unpack(vim.api.nvim_win_get_cursor(0))
+  local ns = ts_utils.get_struct_node_at_pos(row + 1, col + 1)
+  if not ns then
+    u.notify("not node found at cursor position", "error")
+    return
+  end
 
-    if current_node:type() ~= 'type_spec' then
-        u.notify("not struct node under cursor","error")
-        return
-    end
-
-    -- TODO: allow to pass a funcName as an argument
-    -- local struct_name = (fn_args.args ~= nil and fn_args.args ~= "") and fn_args.args:gsub('"', '') or vim.treesitter.get_node_text(current_node:child(0), 1)
-    local struct_name = vim.treesitter.get_node_text(current_node:child(0), 1)
-    local struct_type_node = current_node:child(1):child(1)
-
-    local args = {}
-    local struct_initialization = {}
-    for field in struct_type_node:iter_children() do
-        if field:type() == 'field_declaration' then
-            local field_type_node = field:child(1)
-            local field_name_node = field:child(0)
-            if field_type_node and field_name_node then
-                local field_type = vim.treesitter.get_node_text(field_type_node, 1)
-                local field_name = vim.treesitter.get_node_text(field_name_node, 1)
-                local field_name_lower = string.lower(field_name)
-                table.insert(args, field_name_lower .. " " .. field_type)
-                table.insert(struct_initialization, field_name .. ": " .. field_name_lower)
-            end
+  local parameter_identifiers = {}
+  if ns.struct_properties.parameters_node then
+    for i in ns.struct_properties.parameters_node:iter_children() do
+      if i:type() == "type_parameter_declaration" then
+        for j in i:iter_children() do
+          print("j: ", j:type())
+          if j:type() == "identifier" then
+            table.insert(parameter_identifiers, vim.treesitter.get_node_text(j, 1))
+          end
         end
+      end
     end
+  end
 
-    local constructor_code = string.format(
-        "\nfunc New%s(%s) *%s {\n\treturn &%s{%s}\n}\n",
-        struct_name:gsub("^%l", string.upper), table.concat(args, ', '), struct_name, struct_name, table.concat(struct_initialization, ', ')
-    )
+  local args = {}
+  local fields = {}
+  for i in ns.struct_properties.fields_node:iter_children() do
+    if i:type() == "field_declaration" then
+      table.insert(args, vim.treesitter.get_node_text(i, 1))
 
-    vim.fn.append(current_node:end_() + 1, vim.split(constructor_code, '\n'))
+      for j in i:iter_children() do
+        if j:type() == "field_identifier" then
+          local field_name = vim.treesitter.get_node_text(j, 1)
+          table.insert(fields, field_name .. ": " .. field_name)
+        end
+      end
+    end
+  end
+
+  local struct_type = ns.name
+  local parameters_section = ""
+  if #parameter_identifiers > 0 then
+    struct_type = struct_type .. "[" .. table.concat(parameter_identifiers, ", ") .. "]"
+    parameters_section = vim.treesitter.get_node_text(ns.struct_properties.parameters_node, 1)
+  end
+
+  local constructor_code = string.format(
+    "\nfunc New%s%s(%s) *%s {\n\treturn &%s{%s}\n}\n",
+    ns.name:gsub("^%l", string.upper),
+    parameters_section,
+    table.concat(args, ", "),
+    struct_type,
+    struct_type,
+    table.concat(fields, ", ")
+  )
+
+  vim.fn.append(ns.declaring_node:end_() + 1, vim.split(constructor_code, "\n"))
 end