tia-lsp/src/formatter.lua

-- SCL Formatter - Document formatting for the LSP server
-- Stack-based implementation for proper nested structure handling
local M = {}

-- Default patterns for attribute blocks to collapse
-- Users can extend these via options.collapse_patterns
local DEFAULT_COLLAPSE_PATTERNS = {
  "^%s*{%s*EXTERNAL",
  "^%s*{%s*S7_",
  "^%s*{%s*%w+%s*:=",
}

local function position(line, character)
  return { line = line, character = character }
end

local function range(start_line, start_char, end_line, end_char)
  return {
    start = position(start_line, start_char),
    ["end"] = position(end_line, end_char),
  }
end

-- Block types with their corresponding enders
local BLOCKS = {
  ["IF"] = "END_IF",
  ["FOR"] = "END_FOR",
  ["WHILE"] = "END_WHILE",
  ["REPEAT"] = "END_REPEAT",
  ["CASE"] = "END_CASE",
  ["REGION"] = "END_REGION",
  ["FUNCTION"] = "END_FUNCTION",
  ["FUNCTION_BLOCK"] = "END_FUNCTION_BLOCK",
  ["ORGANIZATION_BLOCK"] = "END_ORGANIZATION_BLOCK",
  ["TYPE"] = "END_TYPE",
  ["STRUCT"] = "END_STRUCT",
  ["VAR"] = "END_VAR",
  ["VAR_INPUT"] = "END_VAR",
  ["VAR_OUTPUT"] = "END_VAR",
  ["VAR_IN_OUT"] = "END_VAR",
  ["VAR_TEMP"] = "END_VAR",
  ["VAR_CONSTANT"] = "END_VAR",
  ["VAR_RETAIN"] = "END_VAR",
  ["VAR_NON_RETAIN"] = "END_VAR",
}

local SAME_LEVEL = {
  ["THEN"] = true,
  ["ELSE"] = true,
  ["ELSIF"] = true,
  ["DO"] = true,
  ["OF"] = true,
  ["BEGIN"] = true,
}

local NEEDS_SEMICOLON = {
  ["END_IF"] = true,
  ["END_FOR"] = true,
  ["END_WHILE"] = true,
  ["END_REPEAT"] = true,
  ["END_CASE"] = true,
  ["END_REGION"] = true,
}

local VAR_KEYWORDS = {
  ["VAR"] = true,
  ["VAR_INPUT"] = true,
  ["VAR_OUTPUT"] = true,
  ["VAR_IN_OUT"] = true,
  ["VAR_TEMP"] = true,
  ["VAR_CONSTANT"] = true,
  ["VAR_RETAIN"] = true,
  ["VAR_NON_RETAIN"] = true,
}

local DECLARATIONS = {
  ["FUNCTION"] = true,
  ["FUNCTION_BLOCK"] = true,
  ["ORGANIZATION_BLOCK"] = true,
  ["TYPE"] = true,
  ["STRUCT"] = true,
}

local CONTINUATION_OPERATORS = {
  ["AND"] = true,
  ["OR"] = true,
  ["XOR"] = true,
  ["NOT"] = true,
  ["MOD"] = true,
}

function M.format_document(content, options)
  options = options or {}
  local use_spaces = options.insertSpaces or false
  local indent_size = options.tabSize or 1
  local indent_char = use_spaces and string.rep(" ", indent_size) or "\t"
  
  local collapse_attributes = options.collapseAttributes
  if collapse_attributes == nil then
    collapse_attributes = false
  end
  
  local collapse_patterns = options.collapsePatterns or DEFAULT_COLLAPSE_PATTERNS
  if options.extendCollapsePatterns then
    for _, pattern in ipairs(options.extendCollapsePatterns) do
      table.insert(collapse_patterns, pattern)
    end
  end

  local lines = {}
  for line in content:gmatch("([^\n]*)\n") do
    table.insert(lines, line)
  end
  local last_line = content:match("([^\n]+)$")
  if last_line and last_line ~= "" then
    table.insert(lines, last_line)
  end
  if #lines == 0 then
    return { { range = range(0, 0, 0, 0), newText = "" } }
  end

  local stack = {}
  local formatted_lines = {}

  local function get_indent()
    local level = 0
    for _, block in ipairs(stack) do
      if not block.no_indent then
        level = level + 1
      end
    end
    return string.rep(indent_char, level)
  end

  local function is_assignment_continuation(trimmed)
    if not trimmed then return false end
    local first_word = trimmed:match("^(%w+)")
    if first_word and CONTINUATION_OPERATORS[first_word:upper()] then
      return true
    end
    return false
  end

  local function is_fb_call_start(trimmed)
    if not (trimmed:match("^#?%w+") or trimmed:match('^"')) then
      return false
    end
    local paren_pos = trimmed:find("%(")
    if not paren_pos then return false end
    local assign_pos = trimmed:find(":=")
    if assign_pos and assign_pos < paren_pos then return false end
    return true
  end

  local function fb_call_has_closing(line)
    local open_count, close_count = 0, 0
    for i = 1, #line do
      local c = line:sub(i, i)
      if c == "(" then open_count = open_count + 1
      elseif c == ")" then close_count = close_count + 1 end
    end
    return close_count > 0 and close_count >= open_count
  end

  local function format_fb_call_multiline(fb_lines, base_indent, param_indent)
    local result = {}
    local all_text = table.concat(fb_lines, " ")
    
    local fb_name = all_text:match("^(#?[%w_]+)%s*%(")
    if not fb_name then
      for _, l in ipairs(fb_lines) do
        table.insert(result, base_indent .. l)
      end
      return result
    end
    
    local paren_start = all_text:find("%(")
    local paren_end = all_text:find("%);%s*$") or all_text:find("%)%s*$")
    if not paren_start or not paren_end then
      for _, l in ipairs(fb_lines) do
        table.insert(result, base_indent .. l)
      end
      return result
    end
    
    local params_text = all_text:sub(paren_start + 1, paren_end - 1)
    
    local params = {}
    local current_param = ""
    local depth = 0
    for i = 1, #params_text do
      local c = params_text:sub(i, i)
      if c == "(" then
        depth = depth + 1
        current_param = current_param .. c
      elseif c == ")" then
        depth = depth - 1
        current_param = current_param .. c
      elseif c == "," and depth == 0 then
        local trimmed_param = current_param:match("^%s*(.-)%s*$")
        if trimmed_param ~= "" then
          table.insert(params, trimmed_param)
        end
        current_param = ""
      else
        current_param = current_param .. c
      end
    end
    local last_param = current_param:match("^%s*(.-)%s*$")
    if last_param and last_param ~= "" then
      table.insert(params, last_param)
    end
    
    if #params == 0 then
      return { base_indent .. fb_name .. "();" }
    end
    
    table.insert(result, base_indent .. fb_name .. "(" .. params[1] .. ",")
    
    for i = 2, #params do
      local closing = (i == #params) and ");" or ","
      table.insert(result, param_indent .. params[i] .. closing)
    end
    
    return result
  end

  local function is_case_label(trimmed)
    return trimmed:match("^[%w_#]+%s*:%s*$") ~= nil
  end

  local function extract_variable_name(line)
    local var_name = line:match("^%s*([%w_]+)%s*[%({:]")
    return var_name
  end

  local function collapse_attribute_block(line)
    if not collapse_attributes then
      return line
    end

    local attr_start, attr_end = line:find("{.-}")
    if not attr_start then
      return line
    end

    local before_attr = line:sub(1, attr_start - 1)
    local attr_content = line:sub(attr_start, attr_end)
    local after_attr = line:sub(attr_end + 1)

    local var_name = extract_variable_name(line)

    if options.collapseVariableRules and var_name then
      for _, rule in ipairs(options.collapseVariableRules) do
        local var_pattern = rule.variablePattern
        local attr_pattern = rule.attributePattern
        local should_collapse = rule.collapse

        local var_matches = not var_pattern or var_name:match(var_pattern)
        local attr_matches = not attr_pattern or attr_content:match(attr_pattern)

        if var_matches and attr_matches then
          if should_collapse then
            return before_attr .. "{...}" .. after_attr
          else
            return line
          end
        end
      end
    end

    for _, pattern in ipairs(collapse_patterns) do
      if attr_content:match(pattern) then
        return before_attr .. "{...}" .. after_attr
      end
    end

    return line
  end

  local function get_keyword(trimmed)
    return trimmed:match("^([%w_]+)") or trimmed:match("^(%w+)")
  end

  local function find_case_block()
    for i = #stack, 1, -1 do
      if stack[i].type == "CASE" then
        return stack[i], i
      end
    end
    return nil, nil
  end

  local function pop_to_ender(ender_type)
    while #stack > 0 do
      local top = stack[#stack]
      table.remove(stack)
      local expected_ender = BLOCKS[top.type]
      if expected_ender == ender_type then
        return true
      end
    end
    return false
  end

  -- State tracking
  local in_fb_call = false
  local fb_call_lines = {}
  local fb_base_indent = ""
  local fb_param_indent = ""
  
  local in_assignment = false
  local assignment_continuation_indent = ""

  for line_num, line in ipairs(lines) do
    local trimmed = line:gsub("^%s+", ""):gsub("%s+$", "")
    local kw = get_keyword(trimmed)

    if trimmed == "" then
      table.insert(formatted_lines, "")
      goto continue
    end

    if trimmed:match("^//") or trimmed:match("^%(") then
      local collapsed = collapse_attribute_block(trimmed)
      table.insert(formatted_lines, get_indent() .. collapsed)
      goto continue
    end

    trimmed = collapse_attribute_block(trimmed)

    -- Check if continuing assignment
    if in_assignment and is_assignment_continuation(trimmed) then
      table.insert(formatted_lines, assignment_continuation_indent .. trimmed)
      goto continue
    end
    in_assignment = false

    -- Check if starting FB call
    if not in_fb_call and is_fb_call_start(trimmed) then
      in_fb_call = true
      fb_call_lines = { trimmed }
      fb_base_indent = get_indent()
      local fb_name = trimmed:match("^(#?[%w_]+)")
      if fb_name then
        fb_param_indent = fb_base_indent .. string.rep(" ", #fb_name + 1)
      else
        fb_param_indent = fb_base_indent .. string.rep(" ", 4)
      end
      
      if fb_call_has_closing(trimmed) then
        local formatted_fb = format_fb_call_multiline(fb_call_lines, fb_base_indent, fb_param_indent)
        for _, fline in ipairs(formatted_fb) do
          table.insert(formatted_lines, fline)
        end
        in_fb_call = false
        fb_call_lines = {}
      end
      goto continue
    end

    -- Continue collecting FB call lines
    if in_fb_call then
      table.insert(fb_call_lines, trimmed)
      if fb_call_has_closing(trimmed) then
        local formatted_fb = format_fb_call_multiline(fb_call_lines, fb_base_indent, fb_param_indent)
        for _, fline in ipairs(formatted_fb) do
          table.insert(formatted_lines, fline)
        end
        in_fb_call = false
        fb_call_lines = {}
      end
      goto continue
    end

    if DECLARATIONS[kw] then
      stack = {}
      table.insert(formatted_lines, trimmed)
      table.insert(stack, { type = kw, indent = 0, is_declaration = true })
      goto continue
    end

    if kw and not BLOCKS[kw] then
      for starter, ender in pairs(BLOCKS) do
        if DECLARATIONS[starter] and kw == ender then
          stack = {}
          table.insert(formatted_lines, trimmed)
          goto continue
        end
      end
    end

    if kw == "BEGIN" then
      while #stack > 0 and VAR_KEYWORDS[stack[#stack].type] do
        table.remove(stack)
      end
      table.insert(formatted_lines, get_indent() .. trimmed)
      goto continue
    end

    if VAR_KEYWORDS[kw] then
      while #stack > 0 and not stack[#stack].is_declaration do
        table.remove(stack)
      end
      table.insert(formatted_lines, get_indent() .. trimmed)
      table.insert(stack, { type = kw, indent = #stack, is_var = true })
      goto continue
    end

    if kw == "END_VAR" then
      while #stack > 0 and VAR_KEYWORDS[stack[#stack].type] do
        table.remove(stack)
      end
      table.insert(formatted_lines, get_indent() .. trimmed)
      goto continue
    end

    if kw and BLOCKS[kw] == nil then
      for starter, ender in pairs(BLOCKS) do
        if kw == ender and not DECLARATIONS[starter] then
          pop_to_ender(kw)
          local formatted = get_indent() .. trimmed
          if NEEDS_SEMICOLON[kw] and trimmed:sub(-1) ~= ";" then
            formatted = formatted .. ";"
          end
          table.insert(formatted_lines, formatted)
          goto continue
        end
      end
    end

    if SAME_LEVEL[kw] then
      local formatted
      if kw == "THEN" then
        local if_block = nil
        for i = #stack, 1, -1 do
          if stack[i].type == "IF" then
            if_block = stack[i]
            break
          end
        end
        if if_block then
          formatted = string.rep(indent_char, if_block.indent) .. trimmed
        else
          formatted = get_indent() .. trimmed
        end
        table.insert(formatted_lines, formatted)
        table.insert(stack, { type = "THEN_CONTENT", indent = #stack, no_indent = true })
      elseif kw == "ELSE" or kw == "ELSIF" then
        if #stack > 0 and stack[#stack].type == "THEN_CONTENT" then
          table.remove(stack)
        end
        local if_block = nil
        for i = #stack, 1, -1 do
          if stack[i].type == "IF" then
            if_block = stack[i]
            break
          end
        end
        if if_block then
          formatted = string.rep(indent_char, if_block.indent) .. trimmed
        else
          formatted = get_indent() .. trimmed
        end
        table.insert(formatted_lines, formatted)
        table.insert(stack, { type = "THEN_CONTENT", indent = #stack, no_indent = true })
      else
        formatted = get_indent() .. trimmed
        table.insert(formatted_lines, formatted)
      end
      goto continue
    end

    if is_case_label(trimmed) then
      local case_block, case_idx = find_case_block()
      if case_block then
        while #stack > case_idx do
          table.remove(stack)
        end
        local label_indent = string.rep(indent_char, case_block.indent + 1)
        table.insert(formatted_lines, label_indent .. trimmed)
        table.insert(stack, { type = "CASE_LABEL", indent = case_block.indent + 1, no_indent = true })
      else
        table.insert(formatted_lines, get_indent() .. trimmed)
      end
      goto continue
    end

    if kw and BLOCKS[kw] and not VAR_KEYWORDS[kw] and not DECLARATIONS[kw] then
      local current_level = 0
      for _, block in ipairs(stack) do
        if not block.no_indent then
          current_level = current_level + 1
        end
      end
      local formatted = string.rep(indent_char, current_level) .. trimmed
      table.insert(formatted_lines, formatted)
      table.insert(stack, { type = kw, indent = current_level })
      goto continue
    end

    -- Regular line - check for assignment start
    if trimmed:match(":=") then
      local base_indent = get_indent()
      local assign_pos = trimmed:find(":=")
      if assign_pos then
        local var_part = trimmed:sub(1, assign_pos - 1)
        local var_trimmed = var_part:match("^(.-)%s*$")
        assignment_continuation_indent = base_indent .. string.rep(" ", #var_trimmed + 4)
        
        if not trimmed:match(";%s*$") then
          in_assignment = true
        end
      end
    else
      in_assignment = false
    end

    local formatted = get_indent() .. trimmed
    if NEEDS_SEMICOLON[kw] and trimmed:sub(-1) ~= ";" then
      formatted = formatted .. ";"
    end
    table.insert(formatted_lines, formatted)

    ::continue::
  end

  local result = table.concat(formatted_lines, "\n")
  if content:sub(-1) == "\n" then
    result = result .. "\n"
  end

  return {
    {
      range = range(0, 0, #lines, 0),
      newText = result,
    },
  }
end

function M.get_formatting_options()
  return {
    insertSpaces = false,
    tabSize = 1,
    trimTrailingWhitespace = true,
    insertFinalNewline = true,
    trimFinalNewlines = true,
    collapseAttributes = false,
  }
end

return M