package commands

import (
	"strings"
	"testing"

	"c65gm/internal/compiler"
	"c65gm/internal/preproc"
)

func TestForBasicTO(t *testing.T) {
	tests := []struct {
		name      string
		forLine   string
		setupVars func(*compiler.SymbolTable)
		wantFor   []string
		wantNext  []string
	}{
		{
			name:    "byte var TO byte literal",
			forLine: "FOR i = 0 TO 10",
			setupVars: func(st *compiler.SymbolTable) {
				st.AddVar("i", "", compiler.KindByte, 0)
			},
			wantFor: []string{
				"\tlda #$00",
				"\tsta i",
				"_LOOPSTART1",
				"\tlda #$0a",
				"\tcmp i",
				"\tbcc _LOOPEND1",
			},
			wantNext: []string{
				"\tinc i",
				"\tjmp _LOOPSTART1",
				"_LOOPEND1",
			},
		},
		{
			name:    "word var TO word literal",
			forLine: "FOR counter = 0 TO 1000",
			setupVars: func(st *compiler.SymbolTable) {
				st.AddVar("counter", "", compiler.KindWord, 0)
			},
			wantFor: []string{
				"\tlda #$00",
				"\tsta counter",
				"\tsta counter+1",
				"_LOOPSTART1",
				"\tlda #$03",
				"\tcmp counter+1",
				"\tbcc _LOOPEND1",
				"\tbne _L1",
				"\tlda #$e8",
				"\tcmp counter",
				"\tbcc _LOOPEND1",
				"_L1",
			},
			wantNext: []string{
				"\tinc counter",
				"\tbne _L2",
				"\tinc counter+1",
				"_L2",
				"\tjmp _LOOPSTART1",
				"_LOOPEND1",
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			pragma := preproc.NewPragma()
			ctx := compiler.NewCompilerContext(pragma)
			tt.setupVars(ctx.SymbolTable)

			forCmd := &ForCommand{}
			nextCmd := &NextCommand{}

			forLine := preproc.Line{
				Text:           tt.forLine,
				Kind:           preproc.Source,
				PragmaSetIndex: pragma.GetCurrentPragmaSetIndex(),
			}
			nextLine := preproc.Line{
				Text:           "NEXT",
				Kind:           preproc.Source,
				PragmaSetIndex: pragma.GetCurrentPragmaSetIndex(),
			}

			if err := forCmd.Interpret(forLine, ctx); err != nil {
				t.Fatalf("FOR Interpret() error = %v", err)
			}

			forAsm, err := forCmd.Generate(ctx)
			if err != nil {
				t.Fatalf("FOR Generate() error = %v", err)
			}

			if err := nextCmd.Interpret(nextLine, ctx); err != nil {
				t.Fatalf("NEXT Interpret() error = %v", err)
			}

			nextAsm, err := nextCmd.Generate(ctx)
			if err != nil {
				t.Fatalf("NEXT Generate() error = %v", err)
			}

			if !equalAsm(forAsm, tt.wantFor) {
				t.Errorf("FOR Generate() mismatch\ngot:\n%s\nwant:\n%s",
					strings.Join(forAsm, "\n"),
					strings.Join(tt.wantFor, "\n"))
			}
			if !equalAsm(nextAsm, tt.wantNext) {
				t.Errorf("NEXT Generate() mismatch\ngot:\n%s\nwant:\n%s",
					strings.Join(nextAsm, "\n"),
					strings.Join(tt.wantNext, "\n"))
			}
		})
	}
}

func TestForWithSTEP(t *testing.T) {
	tests := []struct {
		name         string
		forLine      string
		setupVars    func(*compiler.SymbolTable)
		checkNextAsm func([]string) bool
		description  string
	}{
		{
			name:    "byte var TO with STEP 2",
			forLine: "FOR i = 0 TO 10 STEP 2",
			setupVars: func(st *compiler.SymbolTable) {
				st.AddVar("i", "", compiler.KindByte, 0)
			},
			checkNextAsm: func(asm []string) bool {
				// Should contain adc #$02
				for _, line := range asm {
					if strings.Contains(line, "adc #$02") {
						return true
					}
				}
				return false
			},
			description: "STEP 2 should use adc #$02",
		},
		{
			name:    "byte var TO with variable STEP",
			forLine: "FOR i = 0 TO 10 STEP stepval",
			setupVars: func(st *compiler.SymbolTable) {
				st.AddVar("i", "", compiler.KindByte, 0)
				st.AddVar("stepval", "", compiler.KindByte, 0)
			},
			checkNextAsm: func(asm []string) bool {
				// Should contain adc stepval
				for _, line := range asm {
					if strings.Contains(line, "adc stepval") {
						return true
					}
				}
				return false
			},
			description: "variable STEP should use adc variable",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			pragma := preproc.NewPragma()
			ctx := compiler.NewCompilerContext(pragma)
			tt.setupVars(ctx.SymbolTable)

			forCmd := &ForCommand{}
			nextCmd := &NextCommand{}

			forLine := preproc.Line{
				Text:           tt.forLine,
				Kind:           preproc.Source,
				PragmaSetIndex: pragma.GetCurrentPragmaSetIndex(),
			}
			nextLine := preproc.Line{
				Text:           "NEXT",
				Kind:           preproc.Source,
				PragmaSetIndex: pragma.GetCurrentPragmaSetIndex(),
			}

			if err := forCmd.Interpret(forLine, ctx); err != nil {
				t.Fatalf("FOR Interpret() error = %v", err)
			}

			if _, err := forCmd.Generate(ctx); err != nil {
				t.Fatalf("FOR Generate() error = %v", err)
			}

			if err := nextCmd.Interpret(nextLine, ctx); err != nil {
				t.Fatalf("NEXT Interpret() error = %v", err)
			}

			nextAsm, err := nextCmd.Generate(ctx)
			if err != nil {
				t.Fatalf("NEXT Generate() error = %v", err)
			}

			if !tt.checkNextAsm(nextAsm) {
				t.Errorf("%s\ngot:\n%s", tt.description, strings.Join(nextAsm, "\n"))
			}
		})
	}
}

func TestForBreak(t *testing.T) {
	pragma := preproc.NewPragma()
	ctx := compiler.NewCompilerContext(pragma)
	ctx.SymbolTable.AddVar("i", "", compiler.KindByte, 0)

	forCmd := &ForCommand{}
	breakCmd := &BreakCommand{}
	nextCmd := &NextCommand{}

	pragmaIdx := pragma.GetCurrentPragmaSetIndex()

	forLine := preproc.Line{Text: "FOR i = 0 TO 10", Kind: preproc.Source, PragmaSetIndex: pragmaIdx}
	breakLine := preproc.Line{Text: "BREAK", Kind: preproc.Source, PragmaSetIndex: pragmaIdx}
	nextLine := preproc.Line{Text: "NEXT", Kind: preproc.Source, PragmaSetIndex: pragmaIdx}

	if err := forCmd.Interpret(forLine, ctx); err != nil {
		t.Fatalf("FOR Interpret() error = %v", err)
	}

	forAsm, _ := forCmd.Generate(ctx)
	_ = forAsm // body would go here

	if err := breakCmd.Interpret(breakLine, ctx); err != nil {
		t.Fatalf("BREAK Interpret() error = %v", err)
	}

	breakAsm, err := breakCmd.Generate(ctx)
	if err != nil {
		t.Fatalf("BREAK Generate() error = %v", err)
	}

	if err := nextCmd.Interpret(nextLine, ctx); err != nil {
		t.Fatalf("NEXT Interpret() error = %v", err)
	}

	if len(breakAsm) != 1 || !strings.Contains(breakAsm[0], "jmp _LOOPEND") {
		t.Errorf("BREAK should jump to loop end label, got: %v", breakAsm)
	}
}

func TestForNested(t *testing.T) {
	pragma := preproc.NewPragma()
	ctx := compiler.NewCompilerContext(pragma)
	ctx.SymbolTable.AddVar("i", "", compiler.KindByte, 0)
	ctx.SymbolTable.AddVar("j", "", compiler.KindByte, 0)

	pragmaIdx := pragma.GetCurrentPragmaSetIndex()

	for1 := &ForCommand{}
	for2 := &ForCommand{}
	next1 := &NextCommand{}
	next2 := &NextCommand{}

	if err := for1.Interpret(preproc.Line{Text: "FOR i = 0 TO 10", Kind: preproc.Source, PragmaSetIndex: pragmaIdx}, ctx); err != nil {
		t.Fatalf("FOR 1 error = %v", err)
	}
	asm1, err := for1.Generate(ctx)
	if err != nil {
		t.Fatalf("FOR 1 Generate error = %v", err)
	}

	if err := for2.Interpret(preproc.Line{Text: "FOR j = 0 TO 5", Kind: preproc.Source, PragmaSetIndex: pragmaIdx}, ctx); err != nil {
		t.Fatalf("FOR 2 error = %v", err)
	}
	asm2, err := for2.Generate(ctx)
	if err != nil {
		t.Fatalf("FOR 2 Generate error = %v", err)
	}

	if err := next2.Interpret(preproc.Line{Text: "NEXT", Kind: preproc.Source, PragmaSetIndex: pragmaIdx}, ctx); err != nil {
		t.Fatalf("NEXT 2 error = %v", err)
	}
	if err := next1.Interpret(preproc.Line{Text: "NEXT", Kind: preproc.Source, PragmaSetIndex: pragmaIdx}, ctx); err != nil {
		t.Fatalf("NEXT 1 error = %v", err)
	}

	// Find loop start labels in the generated assembly
	loopLabel1 := ""
	loopLabel2 := ""
	for _, line := range asm1 {
		if strings.HasPrefix(line, "_LOOPSTART") {
			loopLabel1 = line
			break
		}
	}
	for _, line := range asm2 {
		if strings.HasPrefix(line, "_LOOPSTART") {
			loopLabel2 = line
			break
		}
	}

	if loopLabel1 == "" || loopLabel2 == "" {
		t.Fatal("Could not find loop labels")
	}
	if loopLabel1 == loopLabel2 {
		t.Error("Nested loops should have different labels")
	}
}

func TestForMixedWithWhile(t *testing.T) {
	pragma := preproc.NewPragma()
	ctx := compiler.NewCompilerContext(pragma)
	ctx.SymbolTable.AddVar("i", "", compiler.KindByte, 0)
	ctx.SymbolTable.AddVar("x", "", compiler.KindByte, 0)

	pragmaIdx := pragma.GetCurrentPragmaSetIndex()

	forCmd := &ForCommand{}
	whileCmd := &WhileCommand{}
	wendCmd := &WendCommand{}
	nextCmd := &NextCommand{}

	// FOR i = 0 TO 10
	//   WHILE x < 5
	//   WEND
	// NEXT

	if err := forCmd.Interpret(preproc.Line{Text: "FOR i = 0 TO 10", Kind: preproc.Source, PragmaSetIndex: pragmaIdx}, ctx); err != nil {
		t.Fatalf("FOR error = %v", err)
	}
	_, _ = forCmd.Generate(ctx)

	if err := whileCmd.Interpret(preproc.Line{Text: "WHILE x < 5", Kind: preproc.Source, PragmaSetIndex: pragmaIdx}, ctx); err != nil {
		t.Fatalf("WHILE error = %v", err)
	}
	_, _ = whileCmd.Generate(ctx)

	if err := wendCmd.Interpret(preproc.Line{Text: "WEND", Kind: preproc.Source, PragmaSetIndex: pragmaIdx}, ctx); err != nil {
		t.Fatalf("WEND error = %v", err)
	}

	if err := nextCmd.Interpret(preproc.Line{Text: "NEXT", Kind: preproc.Source, PragmaSetIndex: pragmaIdx}, ctx); err != nil {
		t.Fatalf("NEXT error = %v", err)
	}
}

func TestForIllegalNesting(t *testing.T) {
	pragma := preproc.NewPragma()
	ctx := compiler.NewCompilerContext(pragma)
	ctx.SymbolTable.AddVar("i", "", compiler.KindByte, 0)
	ctx.SymbolTable.AddVar("x", "", compiler.KindByte, 0)

	pragmaIdx := pragma.GetCurrentPragmaSetIndex()

	forCmd := &ForCommand{}
	whileCmd := &WhileCommand{}
	nextCmd := &NextCommand{}

	// FOR i = 0 TO 10
	//   WHILE x < 5
	//   NEXT  <- ERROR: crossing loop boundaries
	//   WEND

	if err := forCmd.Interpret(preproc.Line{Text: "FOR i = 0 TO 10", Kind: preproc.Source, PragmaSetIndex: pragmaIdx}, ctx); err != nil {
		t.Fatalf("FOR error = %v", err)
	}
	_, _ = forCmd.Generate(ctx)

	if err := whileCmd.Interpret(preproc.Line{Text: "WHILE x < 5", Kind: preproc.Source, PragmaSetIndex: pragmaIdx}, ctx); err != nil {
		t.Fatalf("WHILE error = %v", err)
	}
	_, _ = whileCmd.Generate(ctx)

	// NEXT should fail because of stack mismatch
	err := nextCmd.Interpret(preproc.Line{Text: "NEXT", Kind: preproc.Source, PragmaSetIndex: pragmaIdx}, ctx)
	if err == nil {
		t.Fatal("NEXT should fail when crossing loop boundaries")
	}
	if !strings.Contains(err.Error(), "mismatch") {
		t.Errorf("Wrong error message: %v", err)
	}
}

func TestNextWithoutFor(t *testing.T) {
	pragma := preproc.NewPragma()
	ctx := compiler.NewCompilerContext(pragma)

	cmd := &NextCommand{}
	line := preproc.Line{
		Text:           "NEXT",
		Kind:           preproc.Source,
		PragmaSetIndex: pragma.GetCurrentPragmaSetIndex(),
	}

	err := cmd.Interpret(line, ctx)
	if err == nil {
		t.Fatal("NEXT outside FOR loop should fail")
	}
	if !strings.Contains(err.Error(), "not inside FOR") {
		t.Errorf("Wrong error message: %v", err)
	}
}

func TestForWrongParamCount(t *testing.T) {
	pragma := preproc.NewPragma()
	ctx := compiler.NewCompilerContext(pragma)
	ctx.SymbolTable.AddVar("i", "", compiler.KindByte, 0)

	tests := []string{
		"FOR i",
		"FOR i = 0",
		"FOR i = 0 TO",
		"FOR i = 0 TO 10 STEP",
		"FOR i = 0 TO 10 STEP 2 EXTRA",
	}

	for _, text := range tests {
		cmd := &ForCommand{}
		line := preproc.Line{
			Text:           text,
			Kind:           preproc.Source,
			PragmaSetIndex: pragma.GetCurrentPragmaSetIndex(),
		}

		err := cmd.Interpret(line, ctx)
		if err == nil {
			t.Errorf("Should fail with wrong param count: %s", text)
		}
	}
}

func TestForInvalidDirection(t *testing.T) {
	pragma := preproc.NewPragma()
	ctx := compiler.NewCompilerContext(pragma)
	ctx.SymbolTable.AddVar("i", "", compiler.KindByte, 0)

	cmd := &ForCommand{}
	line := preproc.Line{
		Text:           "FOR i = 0 UPTO 10",
		Kind:           preproc.Source,
		PragmaSetIndex: pragma.GetCurrentPragmaSetIndex(),
	}

	err := cmd.Interpret(line, ctx)
	if err == nil {
		t.Fatal("Should fail with invalid direction keyword")
	}
	if !strings.Contains(err.Error(), "TO") {
		t.Errorf("Wrong error message: %v", err)
	}
}

func TestForDOWNTORejected(t *testing.T) {
	pragma := preproc.NewPragma()
	ctx := compiler.NewCompilerContext(pragma)
	ctx.SymbolTable.AddVar("i", "", compiler.KindByte, 0)

	cmd := &ForCommand{}
	line := preproc.Line{
		Text:           "FOR i = 10 DOWNTO 0",
		Kind:           preproc.Source,
		PragmaSetIndex: pragma.GetCurrentPragmaSetIndex(),
	}

	err := cmd.Interpret(line, ctx)
	if err == nil {
		t.Fatal("Should fail with DOWNTO")
	}
	if !strings.Contains(err.Error(), "not supported") {
		t.Errorf("Wrong error message: %v", err)
	}
}

func TestForZeroStep(t *testing.T) {
	pragma := preproc.NewPragma()
	ctx := compiler.NewCompilerContext(pragma)
	ctx.SymbolTable.AddVar("i", "", compiler.KindByte, 0)

	cmd := &ForCommand{}
	line := preproc.Line{
		Text:           "FOR i = 0 TO 10 STEP 0",
		Kind:           preproc.Source,
		PragmaSetIndex: pragma.GetCurrentPragmaSetIndex(),
	}

	err := cmd.Interpret(line, ctx)
	if err == nil {
		t.Fatal("Should fail with STEP 0")
	}
	if !strings.Contains(err.Error(), "STEP cannot be zero") {
		t.Errorf("Wrong error message: %v", err)
	}
}

func TestForConstVariable(t *testing.T) {
	pragma := preproc.NewPragma()
	ctx := compiler.NewCompilerContext(pragma)
	ctx.SymbolTable.AddConst("LIMIT", "", compiler.KindByte, 10)

	cmd := &ForCommand{}
	line := preproc.Line{
		Text:           "FOR LIMIT = 0 TO 10",
		Kind:           preproc.Source,
		PragmaSetIndex: pragma.GetCurrentPragmaSetIndex(),
	}

	err := cmd.Interpret(line, ctx)
	if err == nil {
		t.Fatal("Should fail when using constant as loop variable")
	}
	if !strings.Contains(err.Error(), "constant") {
		t.Errorf("Wrong error message: %v", err)
	}
}

func TestForUnknownVariable(t *testing.T) {
	pragma := preproc.NewPragma()
	ctx := compiler.NewCompilerContext(pragma)

	cmd := &ForCommand{}
	line := preproc.Line{
		Text:           "FOR unknown = 0 TO 10",
		Kind:           preproc.Source,
		PragmaSetIndex: pragma.GetCurrentPragmaSetIndex(),
	}

	err := cmd.Interpret(line, ctx)
	if err == nil {
		t.Fatal("Should fail with unknown variable")
	}
	if !strings.Contains(err.Error(), "unknown") {
		t.Errorf("Wrong error message: %v", err)
	}
}

func TestForConstantEnd(t *testing.T) {
	pragma := preproc.NewPragma()
	ctx := compiler.NewCompilerContext(pragma)
	ctx.SymbolTable.AddVar("i", "", compiler.KindByte, 0)
	ctx.SymbolTable.AddConst("MAX", "", compiler.KindByte, 100)

	forCmd := &ForCommand{}
	nextCmd := &NextCommand{}

	forLine := preproc.Line{
		Text:           "FOR i = 0 TO MAX",
		Kind:           preproc.Source,
		PragmaSetIndex: pragma.GetCurrentPragmaSetIndex(),
	}
	nextLine := preproc.Line{
		Text:           "NEXT",
		Kind:           preproc.Source,
		PragmaSetIndex: pragma.GetCurrentPragmaSetIndex(),
	}

	if err := forCmd.Interpret(forLine, ctx); err != nil {
		t.Fatalf("FOR Interpret() error = %v", err)
	}

	asm, err := forCmd.Generate(ctx)
	if err != nil {
		t.Fatalf("FOR Generate() error = %v", err)
	}

	found := false
	for _, inst := range asm {
		if strings.Contains(inst, "#$64") { // 100 in hex
			found = true
			break
		}
	}
	if !found {
		t.Error("Constant should be folded to immediate value")
	}

	if err := nextCmd.Interpret(nextLine, ctx); err != nil {
		t.Fatalf("NEXT Interpret() error = %v", err)
	}
}

func TestForWordSTEP(t *testing.T) {
	pragma := preproc.NewPragma()
	ctx := compiler.NewCompilerContext(pragma)
	ctx.SymbolTable.AddVar("counter", "", compiler.KindWord, 0)

	forCmd := &ForCommand{}
	nextCmd := &NextCommand{}

	forLine := preproc.Line{
		Text:           "FOR counter = 0 TO 1000 STEP 256",
		Kind:           preproc.Source,
		PragmaSetIndex: pragma.GetCurrentPragmaSetIndex(),
	}
	nextLine := preproc.Line{
		Text:           "NEXT",
		Kind:           preproc.Source,
		PragmaSetIndex: pragma.GetCurrentPragmaSetIndex(),
	}

	if err := forCmd.Interpret(forLine, ctx); err != nil {
		t.Fatalf("FOR Interpret() error = %v", err)
	}

	if _, err := forCmd.Generate(ctx); err != nil {
		t.Fatalf("FOR Generate() error = %v", err)
	}

	if err := nextCmd.Interpret(nextLine, ctx); err != nil {
		t.Fatalf("NEXT Interpret() error = %v", err)
	}

	nextAsm, err := nextCmd.Generate(ctx)
	if err != nil {
		t.Fatalf("NEXT Generate() error = %v", err)
	}

	// Should handle both low and high bytes
	foundLowAdd := false
	foundHighAdd := false
	for _, inst := range nextAsm {
		if strings.Contains(inst, "adc #$00") {
			foundLowAdd = true
		}
		if strings.Contains(inst, "adc #$01") {
			foundHighAdd = true
		}
	}
	if !foundLowAdd || !foundHighAdd {
		t.Errorf("Word STEP should handle both bytes\ngot:\n%s", strings.Join(nextAsm, "\n"))
	}
}