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)
			},
			// Do-while style: initial guard (constant folded - 0<=10 is true, no code)
			// then loop label
			wantFor: []string{
				"\tlda #$00",
				"\tsta i",
				"_LOOPSTART1",
			},
			// End check before increment
			wantNext: []string{
				"\tlda i",
				"\tcmp #$0a",
				"\tbeq _LOOPEND1",
				"\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)
			},
			// Do-while style: initial guard (constant folded - 0<=1000 is true, no code)
			wantFor: []string{
				"\tlda #$00",
				"\tsta counter",
				"\tsta counter+1",
				"_LOOPSTART1",
			},
			// End check before increment (WORD comparison)
			wantNext: []string{
				"\tlda counter",
				"\tcmp #$e8",
				"\tbne +",
				"\tlda counter+1",
				"\tcmp #$03",
				"\tbeq _LOOPEND1",
				"+",
				"\tinc counter",
				"\tbne _L1",
				"\tinc counter+1",
				"_L1",
				"\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 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 2", // STEP not supported
		"FOR i = 0 TO 10 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 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)
	}

	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)
	}

	// With do-while pattern, the constant end value appears in NEXT's end check
	nextAsm, err := nextCmd.Generate(ctx)
	if err != nil {
		t.Fatalf("NEXT Generate() error = %v", err)
	}

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

func TestForByteMaxEndValue(t *testing.T) {
	// FOR b = 0 TO 255 with BYTE iterator uses do-while pattern:
	// Before incrementing, check if b == end and exit if so.
	// This naturally handles the max value case (255) without overflow.
	pragma := preproc.NewPragma()
	ctx := compiler.NewCompilerContext(pragma)
	ctx.SymbolTable.AddVar("b", "", compiler.KindByte, 0)

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

	forLine := preproc.Line{
		Text:           "FOR b = 0 TO 255",
		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)
	}

	// NEXT should check if b == 255 before incrementing
	// Look for: lda b / cmp #$ff / beq _LOOPEND
	hasLda := false
	hasCmp255 := false
	hasBeq := false
	for _, line := range nextAsm {
		if strings.Contains(line, "lda b") {
			hasLda = true
		}
		if strings.Contains(line, "cmp #$ff") {
			hasCmp255 = true
		}
		if strings.Contains(line, "beq _LOOPEND") {
			hasBeq = true
		}
	}

	if !hasLda || !hasCmp255 || !hasBeq {
		t.Errorf("NEXT should generate overflow check for BYTE TO 255\ngot:\n%s",
			strings.Join(nextAsm, "\n"))
	}
}

func TestForWordMaxEndValue(t *testing.T) {
	// FOR w = 0 TO 65535 with WORD iterator uses do-while pattern.
	// Naturally handles the max value case (65535) without overflow.
	pragma := preproc.NewPragma()
	ctx := compiler.NewCompilerContext(pragma)
	ctx.SymbolTable.AddVar("w", "", compiler.KindWord, 0)

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

	forLine := preproc.Line{
		Text:           "FOR w = 0 TO 65535",
		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)
	}

	// NEXT should check if w == 65535 ($FFFF) before incrementing
	// Look for comparisons with $ff for both bytes
	hasCmpFF := false
	hasBeq := false
	for _, line := range nextAsm {
		if strings.Contains(line, "cmp #$ff") {
			hasCmpFF = true
		}
		if strings.Contains(line, "beq _LOOPEND") {
			hasBeq = true
		}
	}

	if !hasCmpFF || !hasBeq {
		t.Errorf("NEXT should generate overflow check for WORD TO 65535\ngot:\n%s",
			strings.Join(nextAsm, "\n"))
	}
}