package parser_test

import (
	"strings"
	"testing"

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

func TestBasicScopeTracking(t *testing.T) {
	tests := []struct {
		input       string
		assignments []struct {
			variable      string
			isDeclaration bool
		}
		desc string
	}{
		{
			"x = 5",
			[]struct {
				variable      string
				isDeclaration bool
			}{
				{"x", true}, // first assignment to x in global scope
			},
			"single global declaration",
		},
		{
			`x = 5
x = 10`,
			[]struct {
				variable      string
				isDeclaration bool
			}{
				{"x", true},  // declaration
				{"x", false}, // assignment to existing
			},
			"global declaration then assignment",
		},
		{
			`x = 5
y = 10
x = 15`,
			[]struct {
				variable      string
				isDeclaration bool
			}{
				{"x", true},  // declaration
				{"y", true},  // declaration
				{"x", false}, // assignment to existing
			},
			"multiple variables with reassignment",
		},
	}

	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)

			assignmentCount := 0
			for _, stmt := range program.Statements {
				if assign, ok := stmt.(*parser.Assignment); ok {
					if assignmentCount >= len(tt.assignments) {
						t.Fatalf("more assignments than expected")
					}

					expected := tt.assignments[assignmentCount]

					// Check variable name
					ident, ok := assign.Target.(*parser.Identifier)
					if !ok {
						t.Fatalf("expected Identifier, got %T", assign.Target)
					}

					if ident.Value != expected.variable {
						t.Errorf("assignment %d: expected variable %s, got %s",
							assignmentCount, expected.variable, ident.Value)
					}

					// Check declaration status
					if assign.IsDeclaration != expected.isDeclaration {
						t.Errorf("assignment %d (%s): expected IsDeclaration = %t, got %t",
							assignmentCount, expected.variable, expected.isDeclaration, assign.IsDeclaration)
					}

					assignmentCount++
				}
			}

			if assignmentCount != len(tt.assignments) {
				t.Errorf("expected %d assignments, found %d", len(tt.assignments), assignmentCount)
			}
		})
	}
}

func TestBlockScopeTracking(t *testing.T) {
	input := `x = 5
if true then
	y = 10
	x = 15
	z = 20
end
y = 25
z = 30`

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

	// With function/global scoping, all variables are declared at global level
	expectedAssignments := []struct {
		variable      string
		isDeclaration bool
		location      string
	}{
		{"x", true, "global"},  // x = 5
		{"y", true, "global"},  // y = 10 (first in global, even though in if block)
		{"x", false, "global"}, // x = 15 (exists in global)
		{"z", true, "global"},  // z = 20 (first in global, even though in if block)
		{"y", false, "global"}, // y = 25 (already declared in global)
		{"z", false, "global"}, // z = 30 (already declared in global)
	}

	assignments := extractAssignments(program)

	if len(assignments) != len(expectedAssignments) {
		t.Fatalf("expected %d assignments, got %d", len(expectedAssignments), len(assignments))
	}

	for i, expected := range expectedAssignments {
		assign := assignments[i]
		ident, ok := assign.Target.(*parser.Identifier)
		if !ok {
			t.Fatalf("assignment %d: expected Identifier, got %T", i, assign.Target)
		}

		if ident.Value != expected.variable {
			t.Errorf("assignment %d: expected variable %s, got %s",
				i, expected.variable, ident.Value)
		}

		if assign.IsDeclaration != expected.isDeclaration {
			t.Errorf("assignment %d (%s in %s): expected IsDeclaration = %t, got %t",
				i, expected.variable, expected.location, expected.isDeclaration, assign.IsDeclaration)
		}
	}
}

func TestNestedBlockScopes(t *testing.T) {
	input := `x = 1
if true then
	y = 2
	if true then
		z = 3
		x = 4
		y = 5
	end
	z = 6
end`

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

	// With function/global scoping, all variables are declared at global level
	expectedAssignments := []struct {
		variable      string
		isDeclaration bool
		scope         string
	}{
		{"x", true, "global"},  // x = 1
		{"y", true, "global"},  // y = 2 (first in global)
		{"z", true, "global"},  // z = 3 (first in global)
		{"x", false, "global"}, // x = 4 (exists in global)
		{"y", false, "global"}, // y = 5 (exists in global)
		{"z", false, "global"}, // z = 6 (exists in global)
	}

	assignments := extractAssignments(program)

	if len(assignments) != len(expectedAssignments) {
		t.Fatalf("expected %d assignments, got %d", len(expectedAssignments), len(assignments))
	}

	for i, expected := range expectedAssignments {
		assign := assignments[i]
		ident, ok := assign.Target.(*parser.Identifier)
		if !ok {
			t.Fatalf("assignment %d: expected Identifier, got %T", i, assign.Target)
		}

		if ident.Value != expected.variable {
			t.Errorf("assignment %d: expected variable %s, got %s",
				i, expected.variable, ident.Value)
		}

		if assign.IsDeclaration != expected.isDeclaration {
			t.Errorf("assignment %d (%s in %s): expected IsDeclaration = %t, got %t",
				i, expected.variable, expected.scope, expected.isDeclaration, assign.IsDeclaration)
		}
	}
}

func TestFunctionScopeTracking(t *testing.T) {
	input := `x = 1
callback = fn(a, b)
	c = a + b
	x = 10
	return c
end
c = 20`

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

	// Extract assignments from the program
	assignments := extractAssignments(program)

	expectedAssignments := []struct {
		variable      string
		isDeclaration bool
		scope         string
	}{
		{"x", true, "global"},        // x = 1
		{"callback", true, "global"}, // callback = fn...
		{"c", true, "function"},      // c = a + b (first in function scope)
		{"x", false, "function"},     // x = 10 (exists in global)
		{"c", true, "global"},        // c = 20 (c was local to function)
	}

	if len(assignments) != len(expectedAssignments) {
		t.Fatalf("expected %d assignments, got %d", len(expectedAssignments), len(assignments))
	}

	for i, expected := range expectedAssignments {
		assign := assignments[i]
		ident, ok := assign.Target.(*parser.Identifier)
		if !ok {
			t.Fatalf("assignment %d: expected Identifier, got %T", i, assign.Target)
		}

		if ident.Value != expected.variable {
			t.Errorf("assignment %d: expected variable %s, got %s",
				i, expected.variable, ident.Value)
		}

		if assign.IsDeclaration != expected.isDeclaration {
			t.Errorf("assignment %d (%s in %s): expected IsDeclaration = %t, got %t",
				i, expected.variable, expected.scope, expected.isDeclaration, assign.IsDeclaration)
		}
	}
}

func TestLoopScopeTracking(t *testing.T) {
	tests := []struct {
		input       string
		desc        string
		assignments []struct {
			variable      string
			isDeclaration bool
			scope         string
		}
	}{
		{
			`x = 1
for i = 1, 10 do
	y = i
	x = 5
end
y = 20`,
			"numeric for loop",
			[]struct {
				variable      string
				isDeclaration bool
				scope         string
			}{
				{"x", true, "global"},  // x = 1
				{"y", true, "global"},  // y = i (first in global, even though in loop)
				{"x", false, "global"}, // x = 5 (exists in global)
				{"y", false, "global"}, // y = 20 (already declared in global)
			},
		},
		{
			`arr = {1, 2, 3}
for k, v in arr do
	sum = sum + v
	arr = nil
end
sum = 0`,
			"for-in loop",
			[]struct {
				variable      string
				isDeclaration bool
				scope         string
			}{
				{"arr", true, "global"},  // arr = {1, 2, 3}
				{"sum", true, "global"},  // sum = sum + v (first in global)
				{"arr", false, "global"}, // arr = nil (exists in global)
				{"sum", false, "global"}, // sum = 0 (already declared in global)
			},
		},
		{
			`running = true
while running do
	count = count + 1
	running = false
end
count = 0`,
			"while loop",
			[]struct {
				variable      string
				isDeclaration bool
				scope         string
			}{
				{"running", true, "global"},  // running = true
				{"count", true, "global"},    // count = count + 1 (first in global)
				{"running", false, "global"}, // running = false (exists in global)
				{"count", false, "global"},   // count = 0 (already declared in global)
			},
		},
	}

	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)

			assignments := extractAssignments(program)

			if len(assignments) != len(tt.assignments) {
				t.Fatalf("expected %d assignments, got %d", len(tt.assignments), len(assignments))
			}

			for i, expected := range tt.assignments {
				assign := assignments[i]
				ident, ok := assign.Target.(*parser.Identifier)
				if !ok {
					t.Fatalf("assignment %d: expected Identifier, got %T", i, assign.Target)
				}

				if ident.Value != expected.variable {
					t.Errorf("assignment %d: expected variable %s, got %s",
						i, expected.variable, ident.Value)
				}

				if assign.IsDeclaration != expected.isDeclaration {
					t.Errorf("assignment %d (%s in %s): expected IsDeclaration = %t, got %t",
						i, expected.variable, expected.scope, expected.isDeclaration, assign.IsDeclaration)
				}
			}
		})
	}
}

func TestMemberAccessNotDeclaration(t *testing.T) {
	input := `table = {}
table.key = "value"
table["index"] = 42
arr = {1, 2, 3}
arr[1] = 99`

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

	expectedAssignments := []struct {
		variable       string
		isDeclaration  bool
		isMemberAccess bool
	}{
		{"table", true, false}, // table = {} (declaration)
		{"", false, true},      // table.key = "value" (member access, never declaration)
		{"", false, true},      // table["index"] = 42 (member access, never declaration)
		{"arr", true, false},   // arr = {1, 2, 3} (declaration)
		{"", false, true},      // arr[1] = 99 (member access, never declaration)
	}

	assignmentCount := 0
	for _, stmt := range program.Statements {
		if assign, ok := stmt.(*parser.Assignment); ok {
			if assignmentCount >= len(expectedAssignments) {
				t.Fatalf("more assignments than expected")
			}

			expected := expectedAssignments[assignmentCount]

			if expected.isMemberAccess {
				// Should not be an identifier
				if _, ok := assign.Target.(*parser.Identifier); ok {
					t.Errorf("assignment %d: expected member access, got Identifier", assignmentCount)
				}

				// Member access should never be a declaration
				if assign.IsDeclaration {
					t.Errorf("assignment %d: member access should never be declaration", assignmentCount)
				}
			} else {
				// Should be an identifier
				ident, ok := assign.Target.(*parser.Identifier)
				if !ok {
					t.Errorf("assignment %d: expected Identifier, got %T", assignmentCount, assign.Target)
				} else if ident.Value != expected.variable {
					t.Errorf("assignment %d: expected variable %s, got %s",
						assignmentCount, expected.variable, ident.Value)
				}

				if assign.IsDeclaration != expected.isDeclaration {
					t.Errorf("assignment %d (%s): expected IsDeclaration = %t, got %t",
						assignmentCount, expected.variable, expected.isDeclaration, assign.IsDeclaration)
				}
			}

			assignmentCount++
		}
	}
}

func TestComplexScopeScenario(t *testing.T) {
	input := `global_var = "global"
counter = 0

callback = fn(x)
	local_var = x * 2
	global_var = "modified"
	
	if local_var > 10 then
		temp = "high"
		counter = counter + 1
	else
		temp = "low"
		counter = counter - 1
	end
	
	for i = 1, 3 do
		temp = temp + i
		local_var = local_var + i
	end
	
	return local_var
end

temp = "global_temp"
local_var = "global_local"`

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

	expectedAssignments := []struct {
		variable      string
		isDeclaration bool
		scope         string
	}{
		{"global_var", true, "global"},    // global_var = "global"
		{"counter", true, "global"},       // counter = 0
		{"callback", true, "global"},      // callback = fn...
		{"local_var", true, "function"},   // local_var = x * 2 (first in function)
		{"global_var", false, "function"}, // global_var = "modified" (exists in global)
		{"temp", true, "function"},        // temp = "high" (first in function scope)
		{"counter", false, "function"},    // counter = counter + 1 (exists in global)
		{"temp", false, "function"},       // temp = "low" (already declared in function)
		{"counter", false, "function"},    // counter = counter - 1 (exists in global)
		{"temp", false, "function"},       // temp = temp + i (already declared in function)
		{"local_var", false, "function"},  // local_var = local_var + i (exists in function)
		{"temp", true, "global"},          // temp = "global_temp" (temp was local to function)
		{"local_var", true, "global"},     // local_var = "global_local" (local_var was local to function)
	}

	assignments := extractAssignments(program)

	if len(assignments) != len(expectedAssignments) {
		t.Fatalf("expected %d assignments, got %d", len(expectedAssignments), len(assignments))
	}

	for i, expected := range expectedAssignments {
		assign := assignments[i]
		ident, ok := assign.Target.(*parser.Identifier)
		if !ok {
			t.Fatalf("assignment %d: expected Identifier, got %T", i, assign.Target)
		}

		if ident.Value != expected.variable {
			t.Errorf("assignment %d: expected variable %s, got %s",
				i, expected.variable, ident.Value)
		}

		if assign.IsDeclaration != expected.isDeclaration {
			t.Errorf("assignment %d (%s in %s): expected IsDeclaration = %t, got %t",
				i, expected.variable, expected.scope, expected.isDeclaration, assign.IsDeclaration)
		}
	}
}

func TestScopeStringRepresentation(t *testing.T) {
	input := `x = 5
if true then
	y = 10
	x = 15
end
y = 20`

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

	// With function/global scoping:
	// x = 5: declaration in global
	// y = 10: declaration in global (first occurrence)
	// x = 15: assignment in global (already exists)
	// y = 20: assignment in global (already exists)
	expected := `local x = 5.00
if true then
	local y = 10.00
	x = 15.00
end
y = 20.00`

	result := strings.TrimSpace(program.String())
	if result != expected {
		t.Errorf("expected:\n%s\n\ngot:\n%s", expected, result)
	}
}

// Helper function to extract all assignments from a program recursively
func extractAssignments(program *parser.Program) []*parser.Assignment {
	var assignments []*parser.Assignment

	for _, stmt := range program.Statements {
		assignments = append(assignments, extractAssignmentsFromStatement(stmt)...)
	}

	return assignments
}

func extractAssignmentsFromStatement(stmt parser.Statement) []*parser.Assignment {
	var assignments []*parser.Assignment

	switch s := stmt.(type) {
	case *parser.Assignment:
		assignments = append(assignments, s)

		// Check if the value is a function literal with assignments in body
		if fn, ok := s.Value.(*parser.FunctionLiteral); ok {
			for _, bodyStmt := range fn.Body {
				assignments = append(assignments, extractAssignmentsFromStatement(bodyStmt)...)
			}
		}

	case *parser.IfStatement:
		// Extract from if body
		for _, bodyStmt := range s.Body {
			assignments = append(assignments, extractAssignmentsFromStatement(bodyStmt)...)
		}

		// Extract from elseif bodies
		for _, elseif := range s.ElseIfs {
			for _, bodyStmt := range elseif.Body {
				assignments = append(assignments, extractAssignmentsFromStatement(bodyStmt)...)
			}
		}

		// Extract from else body
		for _, bodyStmt := range s.Else {
			assignments = append(assignments, extractAssignmentsFromStatement(bodyStmt)...)
		}

	case *parser.WhileStatement:
		for _, bodyStmt := range s.Body {
			assignments = append(assignments, extractAssignmentsFromStatement(bodyStmt)...)
		}

	case *parser.ForStatement:
		for _, bodyStmt := range s.Body {
			assignments = append(assignments, extractAssignmentsFromStatement(bodyStmt)...)
		}

	case *parser.ForInStatement:
		for _, bodyStmt := range s.Body {
			assignments = append(assignments, extractAssignmentsFromStatement(bodyStmt)...)
		}
	}

	return assignments
}