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 = {
  -- Match S7 variable attributes: {EXTERNALACCESSIBLE := 'false'; ...}
  "^%s*{%s*EXTERNAL",
  -- Match S7 block attributes: { S7_Optimized_Access := 'TRUE' } or {S7_...}
  "^%s*{%s*S7_",
  -- Match generic attribute blocks with assignments
  "^%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 = {
  -- Control structures
  ["IF"] = "END_IF",
  ["FOR"] = "END_FOR",
  ["WHILE"] = "END_WHILE",
  ["REPEAT"] = "END_REPEAT",
  ["CASE"] = "END_CASE",
  ["REGION"] = "END_REGION",
  -- Declarations  
  ["FUNCTION"] = "END_FUNCTION",
  ["FUNCTION_BLOCK"] = "END_FUNCTION_BLOCK",
  ["ORGANIZATION_BLOCK"] = "END_ORGANIZATION_BLOCK",
  ["TYPE"] = "END_TYPE",
  ["STRUCT"] = "END_STRUCT",
  -- VAR sections (special handling)
  ["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",
}

-- Keywords that don't change indent level themselves (handled specially)
local SAME_LEVEL = {
  ["THEN"] = true,
  ["ELSE"] = true,
  ["ELSIF"] = true,
  ["DO"] = true,
  ["OF"] = true,
  ["BEGIN"] = true,
}

-- Control structure enders that need semicolons
local NEEDS_SEMICOLON = {
  ["END_IF"] = true,
  ["END_FOR"] = true,
  ["END_WHILE"] = true,
  ["END_REPEAT"] = true,
  ["END_CASE"] = true,
  ["END_REGION"] = true,
}

-- VAR keywords for special handling
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,
}

-- Declaration keywords (reset indent)
local DECLARATIONS = {
  ["FUNCTION"] = true,
  ["FUNCTION_BLOCK"] = true,
  ["ORGANIZATION_BLOCK"] = true,
  ["TYPE"] = true,
  ["STRUCT"] = 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"
  
  -- Attribute block collapsing options
  -- NOTE: Disabled by default to preserve compatibility with attr_toggle.
  -- When formatter collapses blocks, the original content is lost.
  -- Use :SCLCollapseAllAttrBlocks after formatting instead.
  local collapse_attributes = options.collapseAttributes
  if collapse_attributes == nil then
    collapse_attributes = false  -- Default: disabled to preserve toggle functionality
  end
  
  -- User can provide custom patterns or extend defaults
  local collapse_patterns = options.collapsePatterns or DEFAULT_COLLAPSE_PATTERNS
  if options.extendCollapsePatterns then
    -- Extend defaults with user patterns
    for _, pattern in ipairs(options.extendCollapsePatterns) do
      table.insert(collapse_patterns, pattern)
    end
  end

  -- Parse content into lines
  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

  -- Stack of active blocks
  -- Each entry: { type = "IF", indent = 1, is_case = false }
  local stack = {}
  local formatted_lines = {}

  -- FB call state
  local in_fb_call = false
  local fb_name_length = 0
  local fb_call_first_line = false
  local fb_paren_offset = 0  -- Offset to align parameters with opening paren

  -- Helper: get current indent string
  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

  -- Helper: get continuation indent for FB calls
  local function get_continuation_indent()
    local base = get_indent()
    return base .. string.rep(" ", fb_paren_offset)
  end

-- Helper: check if line is a case label
	local function is_case_label(trimmed)
		return trimmed:match("^[%w_#]+%s*:%s*$") ~= nil
	end

	-- Helper: extract variable name from a line (before any attribute block)
	local function extract_variable_name(line)
		-- Match pattern: variableName{...} or variableName : type
		-- Variable names can contain letters, numbers, underscores
		local var_name = line:match("^%s*([%w_]+)%s*[%({:]")
		return var_name
	end

	-- Helper: collapse attribute blocks like {EXTERNALACCESSIBLE := 'false'; ...} to {...}
	local function collapse_attribute_block(line)
		if not collapse_attributes then
			return line
		end

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

		-- Extract the content before the attribute block
		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)

		-- Extract variable name for per-variable rules
		local var_name = extract_variable_name(line)

		-- Check per-variable collapse rules first (highest priority)
		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

				-- Check if variable name matches
				local var_matches = not var_pattern or var_name:match(var_pattern)

				-- Check if attribute content matches (if specified)
				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
						-- Explicitly don't collapse this match
						return line
					end
				end
			end
		end

		-- Check if this attribute block matches any global collapse pattern
		for _, pattern in ipairs(collapse_patterns) do
			if attr_content:match(pattern) then
				-- Collapse to {...}
				return before_attr .. "{...}" .. after_attr
			end
		end

		return line
	end

  -- Helper: check if line starts FB call
  local function is_fb_call_start(line)
    local trimmed = line:gsub("^%s+", "")
    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
    local open_count, close_count = 0, 0
    for i = 1, #trimmed do
      local c = trimmed:sub(i, i)
      if c == "(" then open_count = open_count + 1
      elseif c == ")" then close_count = close_count + 1 end
    end
    return open_count > 0 and close_count < open_count
  end

  -- Helper: check if FB call ends on this line
  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
    -- FB call ends when we have closes > opens, not when they're equal
    -- (equal means no parens on this line, which shouldn't end the call)
    return close_count > 0 and close_count >= open_count
  end

  -- Helper: check if line ends with comma (FB call continuation)
  local function is_continuation(line)
    return line:gsub("^%s+", ""):gsub("%s+$", ""):match(".*,$") ~= nil
  end

  -- Helper: extract keyword from line
  local function get_keyword(trimmed)
    -- Match VAR_IN_OUT, VAR_CONSTANT, etc. (words with multiple underscores)
    return trimmed:match("^([%w_]+)") or trimmed:match("^(%w+)")
  end

  -- Helper: find case block in stack
  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

  -- Helper: pop blocks until matching ender found
  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

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

		-- Handle empty lines (preserve as-is)
		if trimmed == "" then
			table.insert(formatted_lines, "")
			goto continue
		end

		-- Handle comments (re-indent with current level)
		if trimmed:match("^//") or trimmed:match("^%(") then
			-- Apply attribute block collapsing to comments too (for consistency)
			local collapsed = collapse_attribute_block(trimmed)
			table.insert(formatted_lines, get_indent() .. collapsed)
			goto continue
		end

		-- Apply attribute block collapsing
		trimmed = collapse_attribute_block(trimmed)

    -- Check for FB call start
    if not in_fb_call and is_fb_call_start(trimmed) then
      local fb_name = trimmed:match("^#?([%w_]+)") or trimmed:match('^"?([%w_]+)"?')
      if fb_name then
        fb_name_length = #fb_name
        in_fb_call = true
        fb_call_first_line = true
        -- Calculate paren offset: find position of opening paren + 1
        -- This aligns continuation lines with the content inside the parens
        local paren_pos = trimmed:find("%(")
        if paren_pos then
          fb_paren_offset = paren_pos + 1  -- +1 to align with content after paren
        else
          fb_paren_offset = fb_name_length + 2
        end
      end
    end

    -- Handle FB call lines
    if in_fb_call then
      local is_first = fb_call_first_line
      local has_closing = fb_call_has_closing(trimmed)
      local is_cont = is_continuation(trimmed)

      local formatted
      if has_closing then
        formatted = (is_first and get_indent() or get_continuation_indent()) .. trimmed
        in_fb_call = false
        fb_call_first_line = false
      elseif is_cont then
        formatted = is_first and (get_indent() .. trimmed) or (get_continuation_indent() .. trimmed)
        fb_call_first_line = false
      else
        formatted = get_continuation_indent() .. trimmed
      end
      table.insert(formatted_lines, formatted)
      goto continue
    end

    -- Handle block declarations (FUNCTION_BLOCK, etc.)
    if DECLARATIONS[kw] then
      -- Clear stack for new block
      stack = {}
      table.insert(formatted_lines, trimmed)
      table.insert(stack, { type = kw, indent = 0, is_declaration = true })
      goto continue
    end

    -- Handle declaration enders
    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

    -- Handle BEGIN (transition to code section)
    if kw == "BEGIN" then
      -- Pop any VAR blocks
      while #stack > 0 and VAR_KEYWORDS[stack[#stack].type] do
        table.remove(stack)
      end
      table.insert(formatted_lines, get_indent() .. trimmed)
      -- BEGIN itself doesn't add indent, but code inside does
      goto continue
    end

    -- Handle VAR section start
    if VAR_KEYWORDS[kw] then
      -- Pop to declaration level
      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

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

    -- Handle block enders (END_IF, END_FOR, etc.)
    if kw and BLOCKS[kw] == nil then
      -- Check if it's an ender
      for starter, ender in pairs(BLOCKS) do
        if kw == ender and not DECLARATIONS[starter] then
          -- Pop matching block
          pop_to_ender(kw)
          -- Add semicolon if needed
          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

    -- Handle same-level keywords (THEN, ELSE, ELSIF)
    if SAME_LEVEL[kw] then
      local formatted
      if kw == "THEN" then
        -- Pop IF block temporarily to get correct indent
        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)
        -- Push content block after THEN
        table.insert(stack, { type = "THEN_CONTENT", indent = #stack, no_indent = true })
      elseif kw == "ELSE" or kw == "ELSIF" then
        -- Pop THEN_CONTENT if present
        if #stack > 0 and stack[#stack].type == "THEN_CONTENT" then
          table.remove(stack)
        end
        -- Find IF block for correct indent
        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)
        -- Push new content block
        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

    -- Handle case labels
    if is_case_label(trimmed) then
      local case_block, case_idx = find_case_block()
      if case_block then
        -- Pop to just after CASE
        while #stack > case_idx do
          table.remove(stack)
        end
        -- Case label at CASE level + 1
        local label_indent = string.rep(indent_char, case_block.indent + 1)
        table.insert(formatted_lines, label_indent .. trimmed)
        -- Push case content level
        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

    -- Handle block starters (IF, FOR, CASE, REGION, etc.)
    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 (including variable declarations inside VAR sections)
    local formatted = get_indent() .. trimmed
    -- Add semicolon if it looks like it needs one
    if NEEDS_SEMICOLON[kw] and trimmed:sub(-1) ~= ";" then
      formatted = formatted .. ";"
    end
    table.insert(formatted_lines, formatted)

    ::continue::
  end

  -- Join lines
  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,  -- Disabled by default to preserve toggle functionality
    -- Set to true to enable automatic collapsing during formatting:
    -- collapseAttributes = true,
    -- collapsePatterns = { ... },  -- Override default collapse patterns
    -- extendCollapsePatterns = { ... },  -- Extend default patterns with custom ones
    -- collapseVariableRules = {  -- Per-variable collapse rules (highest priority)
    --   {
    --     variablePattern = "^stat",  -- Match variable names starting with "stat"
    --     attributePattern = "EXTERNAL",  -- Optional: also match attribute content
    --     collapse = true,  -- true to collapse, false to expand
    --   },
    --   {
    --     variablePattern = "^temp",  -- Match variables starting with "temp"
    --     collapse = false,  -- Never collapse these
    --   },
    -- },
  }
end

return M