package parser_test

import (
	"strings"
	"testing"

	"git.sharkk.net/Sharkk/Mako/parser"
)

func TestNumericForLoop(t *testing.T) {
	tests := []struct {
		input    string
		variable string
		hasStep  bool
		desc     string
	}{
		{"for i = 1, 10 do\necho i\nend", "i", false, "basic numeric loop"},
		{"for j = 0, 5, 2 do\nx = j\nend", "j", true, "numeric loop with step"},
		{"for count = 1, 10, 2 do\necho count\nend", "count", true, "step loop"},
	}

	for _, tt := range tests {
		t.Run(tt.desc, func(t *testing.T) {
			l := parser.NewLexer(tt.input)
			p := parser.NewParser(l)
			program := p.ParseProgram()
			checkParserErrors(t, p)

			if len(program.Statements) != 1 {
				t.Fatalf("expected 1 statement, got %d", len(program.Statements))
			}

			stmt, ok := program.Statements[0].(*parser.ForStatement)
			if !ok {
				t.Fatalf("expected ForStatement, got %T", program.Statements[0])
			}

			if stmt.Variable.Value != tt.variable {
				t.Errorf("expected variable %s, got %s", tt.variable, stmt.Variable.Value)
			}

			if stmt.Start == nil {
				t.Error("expected start expression")
			}

			if stmt.End == nil {
				t.Error("expected end expression")
			}

			if tt.hasStep && stmt.Step == nil {
				t.Error("expected step expression")
			}

			if !tt.hasStep && stmt.Step != nil {
				t.Error("unexpected step expression")
			}

			if len(stmt.Body) == 0 {
				t.Error("expected non-empty body")
			}
		})
	}
}

func TestForInLoop(t *testing.T) {
	tests := []struct {
		input     string
		hasKey    bool
		keyName   string
		valueName string
		desc      string
	}{
		{"for v in arr do\necho v\nend", false, "", "v", "single variable for-in"},
		{"for k, v in table do\necho k\nend", true, "k", "v", "key-value for-in"},
		{"for index, item in list do\necho item\nend", true, "index", "item", "descriptive names"},
	}

	for _, tt := range tests {
		t.Run(tt.desc, func(t *testing.T) {
			l := parser.NewLexer(tt.input)
			p := parser.NewParser(l)
			program := p.ParseProgram()
			checkParserErrors(t, p)

			if len(program.Statements) != 1 {
				t.Fatalf("expected 1 statement, got %d", len(program.Statements))
			}

			stmt, ok := program.Statements[0].(*parser.ForInStatement)
			if !ok {
				t.Fatalf("expected ForInStatement, got %T", program.Statements[0])
			}

			if tt.hasKey {
				if stmt.Key == nil {
					t.Error("expected key variable")
				} else if stmt.Key.Value != tt.keyName {
					t.Errorf("expected key %s, got %s", tt.keyName, stmt.Key.Value)
				}
			} else {
				if stmt.Key != nil {
					t.Error("unexpected key variable")
				}
			}

			if stmt.Value == nil {
				t.Error("expected value variable")
			} else if stmt.Value.Value != tt.valueName {
				t.Errorf("expected value %s, got %s", tt.valueName, stmt.Value.Value)
			}

			if stmt.Iterable == nil {
				t.Error("expected iterable expression")
			}

			if len(stmt.Body) == 0 {
				t.Error("expected non-empty body")
			}
		})
	}
}

func TestNestedForLoops(t *testing.T) {
	input := `for i = 1, 3 do
	for j = 1, 2 do
		echo i + j
	end
	for k, v in table do
		echo v
	end
end`

	l := parser.NewLexer(input)
	p := parser.NewParser(l)
	program := p.ParseProgram()
	checkParserErrors(t, p)

	if len(program.Statements) != 1 {
		t.Fatalf("expected 1 statement, got %d", len(program.Statements))
	}

	outerLoop, ok := program.Statements[0].(*parser.ForStatement)
	if !ok {
		t.Fatalf("expected ForStatement, got %T", program.Statements[0])
	}

	if len(outerLoop.Body) != 2 {
		t.Fatalf("expected 2 body statements, got %d", len(outerLoop.Body))
	}

	// First nested loop: numeric
	innerNumeric, ok := outerLoop.Body[0].(*parser.ForStatement)
	if !ok {
		t.Fatalf("expected nested ForStatement, got %T", outerLoop.Body[0])
	}

	if innerNumeric.Variable.Value != "j" {
		t.Errorf("expected variable j, got %s", innerNumeric.Variable.Value)
	}

	// Second nested loop: for-in
	innerForIn, ok := outerLoop.Body[1].(*parser.ForInStatement)
	if !ok {
		t.Fatalf("expected nested ForInStatement, got %T", outerLoop.Body[1])
	}

	if innerForIn.Key.Value != "k" || innerForIn.Value.Value != "v" {
		t.Error("expected key k and value v in nested for-in loop")
	}
}

func TestForLoopWithComplexExpressions(t *testing.T) {
	input := `for i = start, table.size, step + 1 do
	for k, v in data[index] do
		result[k] = v
	end
end`

	l := parser.NewLexer(input)
	p := parser.NewParser(l)
	program := p.ParseProgram()
	checkParserErrors(t, p)

	if len(program.Statements) != 1 {
		t.Fatalf("expected 1 statement, got %d", len(program.Statements))
	}

	stmt, ok := program.Statements[0].(*parser.ForStatement)
	if !ok {
		t.Fatalf("expected ForStatement, got %T", program.Statements[0])
	}

	// Test that complex expressions are parsed as expressions
	if stmt.Start == nil || stmt.End == nil || stmt.Step == nil {
		t.Error("expected all expressions to be non-nil")
	}

	// Test nested for-in with complex iterable
	if len(stmt.Body) != 1 {
		t.Fatalf("expected 1 body statement, got %d", len(stmt.Body))
	}

	nestedStmt, ok := stmt.Body[0].(*parser.ForInStatement)
	if !ok {
		t.Fatalf("expected nested ForInStatement, got %T", stmt.Body[0])
	}

	if nestedStmt.Iterable == nil {
		t.Error("expected iterable expression")
	}
}

func TestForLoopErrors(t *testing.T) {
	tests := []struct {
		input         string
		expectedError string
		desc          string
	}{
		{"for do end", "expected identifier after 'for'", "missing variable"},
		{"for i do end", "expected '=', ',' or 'in' after for loop variable", "missing assignment or in"},
		{"for i = do end", "expected start expression in for loop", "missing start expression"},
		{"for i = 1 do end", "expected ',' after start expression in for loop", "missing comma"},
		{"for i = 1, do end", "expected end expression in for loop", "missing end expression"},
		{"for i = 1, 10 end", "expected 'do' after for loop header", "missing do"},
		{"for i = 1, 10 do", "expected 'end' to close for loop", "missing end"},
		{"for i, do end", "expected identifier after ',' in for loop", "missing second variable"},
		{"for i, j do end", "expected 'in' in for loop", "missing in keyword"},
		{"for i, j in do end", "expected expression after 'in' in for loop", "missing iterable"},
		{"for i in arr end", "expected 'do' after for loop header", "missing do in for-in"},
		{"for i in arr do", "expected 'end' to close for loop", "missing end in for-in"},
	}

	for _, tt := range tests {
		t.Run(tt.desc, func(t *testing.T) {
			l := parser.NewLexer(tt.input)
			p := parser.NewParser(l)
			p.ParseProgram()

			if !p.HasErrors() {
				t.Fatal("expected parsing errors")
			}

			errors := p.Errors()
			found := false
			for _, err := range errors {
				if strings.Contains(err.Message, tt.expectedError) {
					found = true
					break
				}
			}

			if !found {
				errorMsgs := make([]string, len(errors))
				for i, err := range errors {
					errorMsgs[i] = err.Message
				}
				t.Errorf("expected error containing %q, got %v", tt.expectedError, errorMsgs)
			}
		})
	}
}

func TestForLoopStringRepresentation(t *testing.T) {
	tests := []struct {
		input    string
		contains []string
		desc     string
	}{
		{
			"for i = 1, 10 do\necho i\nend",
			[]string{"for i = 1.00, 10.00 do", "echo i", "end"},
			"numeric for loop",
		},
		{
			"for i = 1, 10, 2 do\nx = i\nend",
			[]string{"for i = 1.00, 10.00, 2.00 do", "x = i", "end"},
			"numeric for loop with step",
		},
		{
			"for v in arr do\necho v\nend",
			[]string{"for v in arr do", "echo v", "end"},
			"single variable for-in",
		},
		{
			"for k, v in table do\necho k\nend",
			[]string{"for k, v in table do", "echo k", "end"},
			"key-value for-in",
		},
	}

	for _, tt := range tests {
		t.Run(tt.desc, func(t *testing.T) {
			l := parser.NewLexer(tt.input)
			p := parser.NewParser(l)
			program := p.ParseProgram()
			checkParserErrors(t, p)

			programStr := program.String()
			for _, contain := range tt.contains {
				if !strings.Contains(programStr, contain) {
					t.Errorf("expected string representation to contain %q, got:\n%s", contain, programStr)
				}
			}
		})
	}
}

func TestForLoopInMixedProgram(t *testing.T) {
	input := `arr = {1, 2, 3}
total = 0
for i = 1, 3 do
	total = total + arr[i]
end
echo total

for k, v in arr do
	echo v
end`

	l := parser.NewLexer(input)
	p := parser.NewParser(l)
	program := p.ParseProgram()
	checkParserErrors(t, p)

	if len(program.Statements) != 5 {
		t.Fatalf("expected 5 statements, got %d", len(program.Statements))
	}

	// First: table assignment
	_, ok := program.Statements[0].(*parser.AssignStatement)
	if !ok {
		t.Fatalf("statement 0: expected AssignStatement, got %T", program.Statements[0])
	}

	// Second: variable assignment
	_, ok = program.Statements[1].(*parser.AssignStatement)
	if !ok {
		t.Fatalf("statement 1: expected AssignStatement, got %T", program.Statements[1])
	}

	// Third: numeric for loop
	forStmt, ok := program.Statements[2].(*parser.ForStatement)
	if !ok {
		t.Fatalf("statement 2: expected ForStatement, got %T", program.Statements[2])
	}

	if len(forStmt.Body) != 1 {
		t.Errorf("expected 1 body statement in for loop, got %d", len(forStmt.Body))
	}

	// Fourth: echo statement
	_, ok = program.Statements[3].(*parser.EchoStatement)
	if !ok {
		t.Fatalf("statement 3: expected EchoStatement, got %T", program.Statements[3])
	}

	// Fifth: for-in loop
	forInStmt, ok := program.Statements[4].(*parser.ForInStatement)
	if !ok {
		t.Fatalf("statement 4: expected ForInStatement, got %T", program.Statements[4])
	}

	if len(forInStmt.Body) != 1 {
		t.Errorf("expected 1 body statement in for-in loop, got %d", len(forInStmt.Body))
	}

	// Check that body contains echo statement
	_, ok = forInStmt.Body[0].(*parser.EchoStatement)
	if !ok {
		t.Fatalf("for-in body: expected EchoStatement, got %T", forInStmt.Body[0])
	}
}

func TestForLoopWithMemberAccess(t *testing.T) {
	input := `for i = obj.start, obj.end, obj.step do
	data[i] = result
	result.items[i] = data[i]
end

for key, val in obj.items do
	val.count = val.count + 1
end`

	l := parser.NewLexer(input)
	p := parser.NewParser(l)
	program := p.ParseProgram()
	checkParserErrors(t, p)

	if len(program.Statements) != 2 {
		t.Fatalf("expected 2 statements, got %d", len(program.Statements))
	}

	// First: numeric for with member access in bounds
	forStmt, ok := program.Statements[0].(*parser.ForStatement)
	if !ok {
		t.Fatalf("statement 0: expected ForStatement, got %T", program.Statements[0])
	}

	// Check bounds are dot expressions
	_, ok = forStmt.Start.(*parser.DotExpression)
	if !ok {
		t.Fatalf("expected DotExpression for start, got %T", forStmt.Start)
	}

	_, ok = forStmt.End.(*parser.DotExpression)
	if !ok {
		t.Fatalf("expected DotExpression for end, got %T", forStmt.End)
	}

	_, ok = forStmt.Step.(*parser.DotExpression)
	if !ok {
		t.Fatalf("expected DotExpression for step, got %T", forStmt.Step)
	}

	// Second: for-in with member access
	forInStmt, ok := program.Statements[1].(*parser.ForInStatement)
	if !ok {
		t.Fatalf("statement 1: expected ForInStatement, got %T", program.Statements[1])
	}

	_, ok = forInStmt.Iterable.(*parser.DotExpression)
	if !ok {
		t.Fatalf("expected DotExpression for iterable, got %T", forInStmt.Iterable)
	}
}

func TestBasicWhileLoop(t *testing.T) {
	tests := []struct {
		input       string
		description string
	}{
		{"while true do\necho \"hello\"\nend", "basic while with boolean"},
		{"while x < 10 do\nx = x + 1\nend", "while with comparison"},
		{"while count do\ncount = count - 1\nend", "while with identifier"},
		{"while table.flag do\necho \"running\"\nend", "while with member access"},
	}

	for _, tt := range tests {
		t.Run(tt.description, func(t *testing.T) {
			l := parser.NewLexer(tt.input)
			p := parser.NewParser(l)
			program := p.ParseProgram()
			checkParserErrors(t, p)

			if len(program.Statements) != 1 {
				t.Fatalf("expected 1 statement, got %d", len(program.Statements))
			}

			stmt, ok := program.Statements[0].(*parser.WhileStatement)
			if !ok {
				t.Fatalf("expected WhileStatement, got %T", program.Statements[0])
			}

			if stmt.Condition == nil {
				t.Error("expected non-nil condition")
			}

			if len(stmt.Body) == 0 {
				t.Error("expected non-empty body")
			}
		})
	}
}

func TestWhileLoopWithComplexConditions(t *testing.T) {
	tests := []struct {
		input       string
		description string
	}{
		{"while x + y < 10 do\necho x\nend", "arithmetic condition"},
		{"while arr[i] != nil do\ni = i + 1\nend", "array access condition"},
		{"while obj.count > 0 do\nobj.count = obj.count - 1\nend", "member access condition"},
		{"while (a + b) * c == d do\necho \"match\"\nend", "grouped expression condition"},
	}

	for _, tt := range tests {
		t.Run(tt.description, func(t *testing.T) {
			l := parser.NewLexer(tt.input)
			p := parser.NewParser(l)
			program := p.ParseProgram()
			checkParserErrors(t, p)

			if len(program.Statements) != 1 {
				t.Fatalf("expected 1 statement, got %d", len(program.Statements))
			}

			stmt, ok := program.Statements[0].(*parser.WhileStatement)
			if !ok {
				t.Fatalf("expected WhileStatement, got %T", program.Statements[0])
			}

			// Condition should be parsed as an expression
			if stmt.Condition == nil {
				t.Error("expected non-nil condition")
			}

			if len(stmt.Body) == 0 {
				t.Error("expected non-empty body")
			}
		})
	}
}

func TestWhileLoopWithComplexBody(t *testing.T) {
	input := `while running do
	echo "processing"
	if data[i] then
		result = result + data[i]
		i = i + 1
	else
		running = false
	end
	for j = 1, 3 do
		echo j
	end
end`

	l := parser.NewLexer(input)
	p := parser.NewParser(l)
	program := p.ParseProgram()
	checkParserErrors(t, p)

	if len(program.Statements) != 1 {
		t.Fatalf("expected 1 statement, got %d", len(program.Statements))
	}

	stmt, ok := program.Statements[0].(*parser.WhileStatement)
	if !ok {
		t.Fatalf("expected WhileStatement, got %T", program.Statements[0])
	}

	if len(stmt.Body) != 3 {
		t.Fatalf("expected 3 body statements, got %d", len(stmt.Body))
	}

	// First: echo statement
	_, ok = stmt.Body[0].(*parser.EchoStatement)
	if !ok {
		t.Fatalf("body[0]: expected EchoStatement, got %T", stmt.Body[0])
	}

	// Second: if statement
	_, ok = stmt.Body[1].(*parser.IfStatement)
	if !ok {
		t.Fatalf("body[1]: expected IfStatement, got %T", stmt.Body[1])
	}

	// Third: for statement
	_, ok = stmt.Body[2].(*parser.ForStatement)
	if !ok {
		t.Fatalf("body[2]: expected ForStatement, got %T", stmt.Body[2])
	}
}

func TestNestedWhileLoops(t *testing.T) {
	input := `while outer do
	while inner do
		echo "nested"
		inner = inner - 1
	end
	outer = outer - 1
end`

	l := parser.NewLexer(input)
	p := parser.NewParser(l)
	program := p.ParseProgram()
	checkParserErrors(t, p)

	if len(program.Statements) != 1 {
		t.Fatalf("expected 1 statement, got %d", len(program.Statements))
	}

	outerWhile, ok := program.Statements[0].(*parser.WhileStatement)
	if !ok {
		t.Fatalf("expected WhileStatement, got %T", program.Statements[0])
	}

	if len(outerWhile.Body) != 2 {
		t.Fatalf("expected 2 body statements, got %d", len(outerWhile.Body))
	}

	// First body statement should be nested while
	innerWhile, ok := outerWhile.Body[0].(*parser.WhileStatement)
	if !ok {
		t.Fatalf("expected nested WhileStatement, got %T", outerWhile.Body[0])
	}

	if len(innerWhile.Body) != 2 {
		t.Fatalf("expected 2 inner body statements, got %d", len(innerWhile.Body))
	}

	// Second body statement should be assignment
	_, ok = outerWhile.Body[1].(*parser.AssignStatement)
	if !ok {
		t.Fatalf("expected AssignStatement, got %T", outerWhile.Body[1])
	}
}

func TestWhileLoopWithMemberAccess(t *testing.T) {
	input := `while obj.running do
	data[index] = data[index] + 1
	index = index + 1
	if index >= obj.size then
		obj.running = false
	end
end`

	l := parser.NewLexer(input)
	p := parser.NewParser(l)
	program := p.ParseProgram()
	checkParserErrors(t, p)

	if len(program.Statements) != 1 {
		t.Fatalf("expected 1 statement, got %d", len(program.Statements))
	}

	stmt, ok := program.Statements[0].(*parser.WhileStatement)
	if !ok {
		t.Fatalf("expected WhileStatement, got %T", program.Statements[0])
	}

	// Check condition is dot expression
	dotExpr, ok := stmt.Condition.(*parser.DotExpression)
	if !ok {
		t.Fatalf("expected DotExpression condition, got %T", stmt.Condition)
	}

	if dotExpr.Key != "running" {
		t.Errorf("expected key 'running', got %s", dotExpr.Key)
	}

	if len(stmt.Body) != 3 {
		t.Fatalf("expected 3 body statements, got %d", len(stmt.Body))
	}

	// First assignment: data[index] = ...
	assign1, ok := stmt.Body[0].(*parser.AssignStatement)
	if !ok {
		t.Fatalf("expected AssignStatement, got %T", stmt.Body[0])
	}

	_, ok = assign1.Name.(*parser.IndexExpression)
	if !ok {
		t.Fatalf("expected IndexExpression for assignment target, got %T", assign1.Name)
	}
}

func TestWhileLoopErrors(t *testing.T) {
	tests := []struct {
		input         string
		expectedError string
		desc          string
	}{
		{"while do end", "expected condition after 'while'", "missing condition"},
		{"while true end", "expected 'do' after while condition", "missing do"},
		{"while true do", "expected 'end' to close while loop", "missing end"},
		{"while + do end", "unexpected operator '+'", "invalid condition"},
		{"while true do x =", "expected expression after assignment operator", "incomplete body"},
	}

	for _, tt := range tests {
		t.Run(tt.desc, func(t *testing.T) {
			l := parser.NewLexer(tt.input)
			p := parser.NewParser(l)
			p.ParseProgram()

			if !p.HasErrors() {
				t.Fatal("expected parsing errors")
			}

			errors := p.Errors()
			found := false
			for _, err := range errors {
				if strings.Contains(err.Message, tt.expectedError) {
					found = true
					break
				}
			}

			if !found {
				errorMsgs := make([]string, len(errors))
				for i, err := range errors {
					errorMsgs[i] = err.Message
				}
				t.Errorf("expected error containing %q, got %v", tt.expectedError, errorMsgs)
			}
		})
	}
}

func TestWhileLoopStringRepresentation(t *testing.T) {
	tests := []struct {
		input    string
		contains []string
		desc     string
	}{
		{
			"while true do\necho \"hello\"\nend",
			[]string{"while true do", "echo \"hello\"", "end"},
			"basic while loop",
		},
		{
			"while x < 10 do\nx = x + 1\nend",
			[]string{"while (x < 10.00) do", "x = (x + 1.00)", "end"},
			"while with condition and assignment",
		},
		{
			"while obj.flag do\necho obj.value\nend",
			[]string{"while obj.flag do", "echo obj.value", "end"},
			"while with member access",
		},
	}

	for _, tt := range tests {
		t.Run(tt.desc, func(t *testing.T) {
			l := parser.NewLexer(tt.input)
			p := parser.NewParser(l)
			program := p.ParseProgram()
			checkParserErrors(t, p)

			programStr := program.String()
			for _, contain := range tt.contains {
				if !strings.Contains(programStr, contain) {
					t.Errorf("expected string representation to contain %q, got:\n%s", contain, programStr)
				}
			}
		})
	}
}

func TestWhileLoopInMixedProgram(t *testing.T) {
	input := `counter = 0
total = 0
arr = {1, 2, 3, 4, 5}

while counter < 5 do
	total = total + arr[counter]
	counter = counter + 1
end

echo total

if total > 10 then
	echo "sum is large"
end`

	l := parser.NewLexer(input)
	p := parser.NewParser(l)
	program := p.ParseProgram()
	checkParserErrors(t, p)

	if len(program.Statements) != 6 {
		t.Fatalf("expected 6 statements, got %d", len(program.Statements))
	}

	// First three: assignments
	for i := 0; i < 3; i++ {
		_, ok := program.Statements[i].(*parser.AssignStatement)
		if !ok {
			t.Fatalf("statement %d: expected AssignStatement, got %T", i, program.Statements[i])
		}
	}

	// Fourth: while loop
	whileStmt, ok := program.Statements[3].(*parser.WhileStatement)
	if !ok {
		t.Fatalf("statement 3: expected WhileStatement, got %T", program.Statements[3])
	}

	if len(whileStmt.Body) != 2 {
		t.Errorf("expected 2 body statements in while loop, got %d", len(whileStmt.Body))
	}

	// Fifth: echo statement
	_, ok = program.Statements[4].(*parser.EchoStatement)
	if !ok {
		t.Fatalf("statement 4: expected EchoStatement, got %T", program.Statements[4])
	}

	// Sixth: if statement
	_, ok = program.Statements[5].(*parser.IfStatement)
	if !ok {
		t.Fatalf("statement 5: expected IfStatement, got %T", program.Statements[5])
	}
}

func TestWhileLoopWithAllLoopTypes(t *testing.T) {
	input := `while active do
	for i = 1, 10 do
		echo i
	end
	
	for k, v in data do
		echo v
	end
	
	while inner do
		inner = inner - 1
	end
	
	active = active - 1
end`

	l := parser.NewLexer(input)
	p := parser.NewParser(l)
	program := p.ParseProgram()
	checkParserErrors(t, p)

	if len(program.Statements) != 1 {
		t.Fatalf("expected 1 statement, got %d", len(program.Statements))
	}

	whileStmt, ok := program.Statements[0].(*parser.WhileStatement)
	if !ok {
		t.Fatalf("expected WhileStatement, got %T", program.Statements[0])
	}

	if len(whileStmt.Body) != 4 {
		t.Fatalf("expected 4 body statements, got %d", len(whileStmt.Body))
	}

	// First: numeric for loop
	_, ok = whileStmt.Body[0].(*parser.ForStatement)
	if !ok {
		t.Fatalf("body[0]: expected ForStatement, got %T", whileStmt.Body[0])
	}

	// Second: for-in loop
	_, ok = whileStmt.Body[1].(*parser.ForInStatement)
	if !ok {
		t.Fatalf("body[1]: expected ForInStatement, got %T", whileStmt.Body[1])
	}

	// Third: nested while loop
	_, ok = whileStmt.Body[2].(*parser.WhileStatement)
	if !ok {
		t.Fatalf("body[2]: expected WhileStatement, got %T", whileStmt.Body[2])
	}

	// Fourth: assignment
	_, ok = whileStmt.Body[3].(*parser.AssignStatement)
	if !ok {
		t.Fatalf("body[3]: expected AssignStatement, got %T", whileStmt.Body[3])
	}
}