local test = dofile("test/test_harness.lua")

package.path = package.path .. ";src/?.lua"

local script_path = debug.getinfo(1, "S").source:gsub("^@", ""):match("(.*/)") or "."
script_path = script_path:gsub("^test/", ""):gsub("/test$", "")
if script_path == "" then script_path = "." end

local formatter = dofile(script_path .. "/src/formatter.lua")

local function test_formatter_basic()
  local content = [[FUNCTION "TestFunc" : Int
VAR_INPUT
   a : Int;
END_VAR
BEGIN

END_FUNCTION
]]
  local result = formatter.format_document(content, {})
  test.assert_not_nil(result, "should return formatting result")
  test.assert(result[1].newText, "should have newText")
end

local function test_formatter_function_block()
  local content = [[FUNCTION_BLOCK "TestFB"
VAR_INPUT
   enable : Bool;
   value : Int;
END_VAR

VAR_OUTPUT
   done : Bool;
END_VAR

BEGIN
   #done := #enable;
END_FUNCTION_BLOCK
]]
  local result = formatter.format_document(content, {})
  test.assert_not_nil(result, "should return formatting result")
  local formatted = result[1].newText
  test.assert(formatted:match("FUNCTION_BLOCK"), "should preserve FUNCTION_BLOCK")
  test.assert(formatted:match("VAR_INPUT"), "should preserve VAR_INPUT")
end

local function test_formatter_if_statement()
  local content = [[IF #x > 5 THEN
#y := 1;
ELSIF #x > 3 THEN
#y := 2;
ELSE
#y := 0;
END_IF;]]
  local result = formatter.format_document(content, {})
  local formatted = result[1].newText
  test.assert(formatted:match("IF"), "should preserve IF")
  test.assert(formatted:match("THEN"), "should preserve THEN")
  test.assert(formatted:match("END_IF"), "should preserve END_IF")
end

local function test_formatter_for_loop()
  local content = [[FOR #i := 0 TO 10 BY 1 DO
#arr[#i] := #i;
END_FOR;]]
  local result = formatter.format_document(content, {})
  local formatted = result[1].newText
  test.assert(formatted:match("FOR"), "should preserve FOR")
  test.assert(formatted:match("END_FOR"), "should preserve END_FOR")
end

local function test_formatter_case_statement()
  local content = [[CASE #x OF
1: #y := 10;
2: #y := 20;
ELSE
#y := 0;
END_CASE;]]
  local result = formatter.format_document(content, {})
  local formatted = result[1].newText
  test.assert(formatted:match("CASE"), "should preserve CASE")
  test.assert(formatted:match("END_CASE"), "should preserve END_CASE")
end

local function test_formatter_comments()
  local content = [[// This is a comment
(* This is a block comment *)
(* Another block comment *)
// Another line comment
#x := 1;]]
  local result = formatter.format_document(content, {})
  local formatted = result[1].newText
  test.assert(formatted:match("This is a comment"), "should preserve line comments")
end

local function test_formatter_multiline_assignment_simple()
  local content = [[#x := 1 AND
2 AND
3;]]
  local result = formatter.format_document(content, {})
  local formatted = result[1].newText
  test.assert(formatted:match("#x := 1 AND"), "should preserve assignment start")
end

local function test_formatter_multiline_assignment_with_hash()
  local content = [[#statWarning := #instTurn.warning OR
#instPusher.warning OR
#instClamp.warning;]]
  local result = formatter.format_document(content, {})
  local formatted = result[1].newText
  test.assert(formatted:match("#statWarning :="), "should preserve assignment start")
  test.assert(formatted:match("#instClamp.warning;$"), "should have continuation")
end

local function test_formatter_multiline_assignment_complex()
  local content = [[#interface.status.main.inBasePosition := #equipmentData.turn.status.homePos AND
#equipmentData.pusher.status.homePos AND
NOT #equipmentData.conveyor.status.forward AND
NOT #equipmentData.conveyor.status.backward;]]
  local result = formatter.format_document(content, {})
  local formatted = result[1].newText
  test.assert(formatted:match("#interface.status.main.inBasePosition :="), "should preserve first line")
  test.assert(formatted:match("NOT #equipmentData"), "should preserve NOT continuation")
end

local function test_formatter_multiline_assignment_ends_properly()
  local content = [[#first := 1 AND
2 AND
3;

#second := 10;]]
  local result = formatter.format_document(content, {})
  local formatted = result[1].newText
  test.assert(formatted:match("#first := 1"), "should preserve first assignment")
  test.assert(formatted:match("#second := 10"), "should start new assignment properly")
end

local function test_formatter_fb_call_single_line()
  local content = [[#instMyFB(enable := #x, done => #y);]]
  local result = formatter.format_document(content, {})
  local formatted = result[1].newText
  test.assert(formatted:match("#instMyFB"), "should preserve FB call")
end

local function test_formatter_fb_call_multiline()
  local content = [[#instMyFB(
enable := #x,
done => #y
);]]
  local result = formatter.format_document(content, {})
  local formatted = result[1].newText
  test.assert(formatted:match("#instMyFB"), "should preserve FB call")
end

local function test_formatter_fb_call_with_long_params()
  local content = [[#instPartCounter(enable := #unitStatus.currentSubMode <> "DRY_RUN",
incPartIn := #equipmentData.partCounter.commands.incPartIn,
incGoodPart := #equipmentData.partCounter.commands.incGoodPart,
resetCntr := #equipmentData.partCounter.commands.resetCntr,
status => #equipmentData.partCounter.status);]]
  local result = formatter.format_document(content, {})
  local formatted = result[1].newText
  test.assert(formatted:match("#instPartCounter"), "should preserve FB call")
  test.assert(formatted:match("status =>"), "should preserve output param")
end

local function test_formatter_region()
  local content = [[REGION MyRegion
// code
#x := 1;
END_REGION;]]
  local result = formatter.format_document(content, {})
  local formatted = result[1].newText
  test.assert(formatted:match("REGION"), "should preserve REGION")
  test.assert(formatted:match("END_REGION"), "should preserve END_REGION")
end

local function test_formatter_struct()
  local content = [[TYPE MyStruct
STRUCT
   x : Int;
   y : Bool;
END_STRUCT;
END_TYPE
]]
  local result = formatter.format_document(content, {})
  local formatted = result[1].newText
  test.assert(formatted:match("TYPE"), "should preserve TYPE")
  test.assert(formatted:match("STRUCT"), "should preserve STRUCT")
end

local function test_formatter_var_sections()
  local content = [[VAR_INPUT
   in1 : Bool;
   in2 : Int;
END_VAR

VAR_OUTPUT
   out1 : Bool;
END_VAR

VAR
   stat1 : Int;
END_VAR
]]
  local result = formatter.format_document(content, {})
  local formatted = result[1].newText
  test.assert(formatted:match("VAR_INPUT"), "should preserve VAR_INPUT")
  test.assert(formatted:match("VAR_OUTPUT"), "should preserve VAR_OUTPUT")
  test.assert(formatted:match("\nVAR\n"), "should preserve VAR")
end

local function test_formatter_attribute_block()
  local content = [[VAR
   stat1{ S7_SetPoint := 'True' } : Int;
   stat2 : Bool;
END_VAR
]]
  local result = formatter.format_document(content, {})
  local formatted = result[1].newText
  test.assert(formatted:match("stat1"), "should preserve variable name")
end

local function test_formatter_while_loop()
  local content = [[WHILE #x < 10 DO
#x := #x + 1;
END_WHILE;]]
  local result = formatter.format_document(content, {})
  local formatted = result[1].newText
  test.assert(formatted:match("WHILE"), "should preserve WHILE")
  test.assert(formatted:match("END_WHILE"), "should preserve END_WHILE")
end

local function test_formatter_repeat_loop()
  local content = [[REPEAT
#x := #x + 1;
UNTIL #x >= 10;]]
  local result = formatter.format_document(content, {})
  local formatted = result[1].newText
  test.assert(formatted:match("REPEAT"), "should preserve REPEAT")
  test.assert(formatted:match("UNTIL"), "should preserve UNTIL")
end

local function test_formatter_assignment_with_string()
  local content = [[#str := "Hello World";]]
  local result = formatter.format_document(content, {})
  local formatted = result[1].newText
  test.assert(formatted:match('"Hello World"'), "should preserve string")
end

local function test_formatter_assignment_with_array()
  local content = [[#arr[0] := 1;
#arr[1] := 2;]]
  local result = formatter.format_document(content, {})
  local formatted = result[1].newText
  test.assert(formatted:match("#arr%[0%]"), "should preserve array access")
end

local function run_formatter_tests()
  test.run_suite("Formatter Tests", function()
    print("\n--- Basic Tests ---")
    test_formatter_basic()
    test_formatter_function_block()
    test_formatter_if_statement()
    test_formatter_for_loop()
    test_formatter_case_statement()
    test_formatter_comments()
    test_formatter_while_loop()
    test_formatter_repeat_loop()
    test_formatter_region()
    test_formatter_struct()
    test_formatter_var_sections()
    test_formatter_attribute_block()
    test_formatter_assignment_with_string()
    test_formatter_assignment_with_array()

    print("\n--- Multiline Assignment Tests ---")
    test_formatter_multiline_assignment_simple()
    test_formatter_multiline_assignment_with_hash()
    test_formatter_multiline_assignment_complex()
    test_formatter_multiline_assignment_ends_properly()

    print("\n--- FB Call Tests ---")
    test_formatter_fb_call_single_line()
    test_formatter_fb_call_multiline()
    test_formatter_fb_call_with_long_params()
  end)
end

if vim then
  vim.api.nvim_create_user_command("TestFormatter", function()
    run_formatter_tests()
  end, {})
end

run_formatter_tests()