local ls = require("luasnip")
local s = ls.snippet
local sn = ls.snippet_node
local isn = ls.indent_snippet_node
local t = ls.text_node
local i = ls.insert_node
local f = ls.function_node
local c = ls.choice_node
local d = ls.dynamic_node
local r = ls.restore_node
local events = require("luasnip.util.events")
local ai = require("luasnip.nodes.absolute_indexer")
local extras = require("luasnip.extras")
local l = extras.lambda
local rep = extras.rep
local p = extras.partial
local m = extras.match
local n = extras.nonempty
local dl = extras.dynamic_lambda
local fmt = require("luasnip.extras.fmt").fmt
local fmta = require("luasnip.extras.fmt").fmta
local conds = require("luasnip.extras.expand_conditions")
local postfix = require("luasnip.extras.postfix").postfix
local types = require("luasnip.util.types")
local parse = require("luasnip.util.parser").parse_snippet
local ms = ls.multi_snippet
local k = require("luasnip.nodes.key_indexer").new_key

local arg_template = [[
  {arg}: {type}
]]

local function pyarg()
  return sn(
    nil,
    fmt(arg_template, {
      arg = i(1, "arg"),
      type = i(2, "Any"),
    })
  )
end

local function pyargs(_, _, _, user_args)
  local choices = {
    sn(nil, i(1)),
    sn(nil, {
      not user_args and t(", ") or t(""),
      d(1, pyarg),
      d(2, pyargs, { user_args = { false } }),
    }),
  }
  -- switch order for first call
  if user_args then
    local fst, snd = unpack(choices)
    choices = { snd, fst }
  end

  return sn(nil, c(1, choices))
end

local def_template = [[
def {fname}({args}) -> {rtype}:
  """
  {docs}
  """
  {final}
]]

local def = s(
  "def",
  fmt(def_template, {
    fname = i(1, "fname"),
    args = d(2, pyargs, nil, { user_args = { true } }),
    rtype = i(3, "None"),
    docs = i(4, "Documentation"),
    final = i(5, "pass"),
  }, { priority = 1001 })
)

local defs_template = [[
def {mname}(self, {args}) -> {rtype}:
  """
  {docs}
  """
  {final}
]]

local defs = s(
  "defs",
  fmt(defs_template, {
    mname = i(1, "mname"),
    args = d(2, pyargs, nil, { user_args = { true } }),
    rtype = i(3, "None"),
    docs = i(4, "Documentation"),
    final = i(5, "pass"),
  }, { priority = 1001 })
)

local enum_template = [[
for {i}, {value} in enumerate({iter}):
    {final}
]]

local dot_enum = postfix(".enum", {
  d(1, function(_, parent)
    return sn(
      1,
      fmt(enum_template, {
        i = i(1, "i"),
        value = i(2, "value"),
        iter = t(parent.env.POSTFIX_MATCH),
        final = i(3, "pass"),
      })
    )
  end),
})

local enum = s(
  "enum",
  fmt(enum_template, {
    i = i(1, "i"),
    value = i(2, "value"),
    iter = i(3, "iter"),
    final = i(4, "pass"),
  })
)

local for_template = [[
for {item} in {iter}:
    {final}
]]

local dot_for = postfix(".for", {
  d(1, function(_, parent)
    return sn(
      1,
      fmt(for_template, {
        item = i(1, "item"),
        iter = t(parent.env.POSTFIX_MATCH),
        final = i(2, "pass"),
      })
    )
  end),
})

local items_template = [[
for {key}, {value} in {iter}:
    {final}
]]

local dot_items = postfix(".items", {
  d(1, function(_, parent)
    return sn(
      1,
      fmt(items_template, {
        key = i(1, "key"),
        value = i(2, "value"),
        iter = t(parent.env.POSTFIX_MATCH),
        final = i(3, "pass"),
      })
    )
  end),
})

local try_template = [[
try:
    {raises}
except {exception} as {ename}:
    {final}
]]

local dot_try = postfix(".try", {
  d(1, function(_, parent)
    return sn(
      1,
      fmt(try_template, {
        raises = t(parent.env.POSTFIX_MATCH),
        exception = i(1, "Exception"),
        ename = i(2, "e"),
        final = i(3, "pass"),
      })
    )
  end),
})

local parr = s(
  "parr",
  fmt(
    [[
    :param {name}: {description}
    :type {name}: {type}
    ]],
    {
      name = i(1, "name"),
      description = i(2, "description"),
      type = i(3, "type"),
    },
    {
      repeat_duplicates = true,
    }
  )
)

local retr = s(
  "retr",
  fmt(
    [[
    :return: {description}
    :rtype: {rtype}
    ]],
    {
      description = i(1, "description"),
      rtype = i(2, "rtype"),
    }
  )
)

local raisr = s(
  "raisr",
  fmt(
    [[
    :raises {exception}: {description}
    ]],
    {
      exception = i(1, "Exception"),
      description = i(2, "description"),
    }
  )
)

return {
  def,
  defs,
  dot_enum,
  dot_for,
  dot_items,
  dot_try,
  enum,
  parr,
  retr,
  raisr,
}