1 files changed, 100 insertions, 55 deletions

diff --git a/lua/snippets/java.lua b/lua/snippets/java.lua
index 1227168..50fef37 100644
--- a/lua/snippets/java.lua
+++ b/lua/snippets/java.lua
@@ -1,92 +1,137 @@
 require('core.snippets.shorthands')
 require('core.snippets.functions')
 
+--- shortcut to choose between java access modifiers
+---@param idx number index of the node
+---@return table choice_node
+local function access_modifiers(idx)
+  return c(idx, {
+    t("public "),
+    t("private "),
+    t("protected "),
+    t("default "),
+  })
+end
+
+--- shortcut to choose between java modifiers
+---@param idx number index of the node
+---@return table choice_node
+local function modifiers(idx)
+  return c(idx, {
+    t(""),
+    t("abstract "),
+    t("final ")
+  })
+end
+
+--- find the current class and return its name if not available defaults to
+--- the file name
+---@return string
+local function java_class()
+  -- set the starting node to the node under the cursor
+  local node = require("nvim-treesitter.ts_utils").get_node_at_cursor()
+
+  while node do
+    -- check if we're in a class
+    if node:type() == "class_declaration" then
+      -- find the class name in the class declaration and return it
+      for i = 0, node:child_count() - 1 do
+        local child = node:child(i)
+        if child and child:type() == "identifier" then
+          return vim.treesitter.get_node_text(child, 0)
+        end
+      end
+    end
+    node = node:parent()
+  end
+
+  -- if no class can be found default to the current file name
+  return file_name()
+end
+
+
 return {
   -- method snippet
-  s("fn", {
-    c(1, {
-      t("public "),
-      t("private "),
-    }),
-    c(2, {
-      t("void"),
-      t("String"),
-      t("char"),
-      t("int"),
-      t("double"),
-      t("boolean"),
-      i(nil, "myType"),
-    }),
-    t(" "),
-    i(3, "myFunc"),
+  s({ trig = [[fn\|constr]], trigEngine = "vim" }, {
+    access_modifiers(1),
+    modifiers(2),
+    d(3, function(_, snip)
+      if snip.trigger == "constr" then
+        return sn(nil, {
+          f(java_class, {})
+        })
+      else
+        return sn(nil, {
+          c(1, {
+            t("void"),
+            t("int"),
+            t("double"),
+            t("boolean"),
+            t("String"),
+            t("char"),
+            i(nil, "myType"),
+          }),
+          t(" "),
+          i(2, "myFunc"),
+        })
+      end
+    end, {}),
     t("("),
     i(4),
     t(")"),
     c(5, {
       t(""),
-      sn(nil, {t({ "", " throws " }), i(1)}),
+      sn(nil, { t({ "", " throws " }), i(1, "Error") }),
     }),
-    t({ " {", "\t" }),
-    i(0),
-    t({ "", "}" }),
+    d(6, function(args)
+      if args[1][1] == "abstract " then
+        -- if the function is abstract do not include a body
+        return sn(nil, {
+          t(";")
+        })
+      else
+        -- if the function isn't abstract include a body
+        return sn(nil, {
+          t({ " {", "\t" }),
+          i(0),
+          t({ "", "}" })
+        })
+      end
+    end, { 2 })
   }),
 
   s("main", {
     t({ "public static void main(String[] args) {", "\t" }),
     i(0),
-    t({ "", "}" }),
-  }),
-
-  -- constructor snippet
-  s("constr", {
-    c(1, {
-      t("public "),
-      t("private "),
-      t("protected ")
-    }),
-    f(file_name, {}),
-    t("("),
-    i(2),
-    t(")"),
-    c(3, {
-      t(""),
-      sn(nil, {t({ "", " throws " }), i(1)}),
-    }),
-    t({ " {", "\t" }),
-    i(0),
-    t({ "", "}" }),
+    t({ "", "}" })
   }),
 
   -- class snippet
   s("class", {
-    c(1, {
-      t("public "),
-      t("private "),
-      t("protected "),
-      t("")
-    }),
-    c(2, {
+    access_modifiers(1),
+    modifiers(2),
+    c(3, {
       t("class "),
       t("interface ")
     }),
-    c(3, {
-      i(0, "MyClass"),
-      f(file_name, {})
-    }),
     c(4, {
+      f(file_name, {}),
+      i(0, "MyClass")
+    }),
+    c(5, {
       t(" "),
       sn(nil, { t({" implements "}), i(1), t(" ") }),
       sn(nil, { t({" extends "}), i(1), t(" ") }),
     }),
     t({ "{", "\t" }),
     i(0),
-    t({ "", "}" }),
+    t({ "", "}" })
   }),
 
-  -- pacakge snippet
+  -- pacakage snippet
   s("package", {
     t("package "),
-    f(function(args, parent, user_args)
+    f(function()
       -- get path
       local dir = vim.fn.expand("%:h")
       -- remove prefix