package optimizer

import (
	"testing"
)

func lines(s ...string) []string { return s }

func TestPassLoadElimination(t *testing.T) {
	tests := []struct {
		name     string
		input    []string
		expected int // expected number of lines after optimization
	}{
		{
			name:     "redundant same variable",
			input:    lines("\tlda x", "\tsta y", "\tlda x"),
			expected: 2, // third lda removed
		},
		{
			name:     "not redundant - store to same variable",
			input:    lines("\tlda x", "\tsta x", "\tlda x"),
			expected: 3, // X was stored, so A is stale
		},
		{
			name:     "not redundant - arithmetic in between",
			input:    lines("\tlda x", "\tadc #1", "\tlda x"),
			expected: 3, // A was modified
		},
		{
			name:     "redundant - label resets state",
			input:    lines("\tlda x", "label", "\tlda x"),
			expected: 3, // label resets state, second lda not redundant
		},
		{
			name:     "x register redundant",
			input:    lines("\tldx counter", "\tldy counter", "\tldx counter"),
			expected: 2, // third ldx is redundant (X still holds counter)
		},
		{
			name:     "immediate loads not affected",
			input:    lines("\tlda #5", "\tsta x", "\tlda #5"),
			expected: 3, // handled by imm pass
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			parsed := parseLines(tt.input)
			result := passLoadElimination(parsed)
			cleaned := stripOptMarkers(result)
			if got := len(cleaned); got != tt.expected {
				t.Errorf("got %d lines, want %d\ninput: %v\noutput: %v", got, tt.expected, tt.input, linesToString(cleaned))
			}
		})
	}
}

func TestPassImmElimination(t *testing.T) {
	tests := []struct {
		name     string
		input    []string
		expected int
	}{
		{
			name:     "redundant same immediate",
			input:    lines("\tlda #$05", "\tsta y", "\tlda #$05"),
			expected: 2, // third removed
		},
		{
			name:     "different immediate values",
			input:    lines("\tlda #5", "\tsta y", "\tlda #3"),
			expected: 3, // different values
		},
		{
			name:     "arithmetic invalidates",
			input:    lines("\tlda #5", "\tadc #1", "\tlda #5"),
			expected: 3, // adc modified A
		},
		{
			name:     "x and y independent",
			input:    lines("\tldx #$00", "\tldy #$00", "\tldx #$00"),
			expected: 2, // first ldx removed? No, let me check: ldx #0 then ldy #0, X still 0, so third redundant
			// Actually: ldx#0(X=0), ldy#0(Y=0,X=0), ldx#0(X=0,redundant)
			// Result: ldx#0, ldy#0 → 2 lines
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			parsed := parseLines(tt.input)
			result := passImmElimination(parsed)
			cleaned := stripOptMarkers(result)
			if got := len(cleaned); got != tt.expected {
				t.Errorf("got %d lines, want %d\ninput: %v\noutput: %v", got, tt.expected, tt.input, linesToString(cleaned))
			}
		})
	}
}

func TestPassJmpNext(t *testing.T) {
	tests := []struct {
		name     string
		input    []string
		expected int
	}{
		{
			name:     "jmp to next label",
			input:    lines("\tjmp _L1", "_L1"),
			expected: 1, // jmp removed
		},
		{
			name:     "jmp not to next label",
			input:    lines("\tjmp _L1", "\tlda x", "_L1"),
			expected: 3, // not immediately followed by label
		},
		{
			name:     "label different from jmp target",
			input:    lines("\tjmp _L1", "_L2"),
			expected: 2, // different labels
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			parsed := parseLines(tt.input)
			result := passJmpNext(parsed)
			cleaned := stripOptMarkers(result)
			if got := len(cleaned); got != tt.expected {
				t.Errorf("got %d lines, want %d\ninput: %v\noutput: %v", got, tt.expected, tt.input, linesToString(cleaned))
			}
		})
	}
}

func TestPassSelfAssignment(t *testing.T) {
	tests := []struct {
		name     string
		input    []string
		expected int
	}{
		{
			name:     "self assignment",
			input:    lines("\tlda x", "\tsta x"),
			expected: 0, // both removed
		},
		{
			name:     "not self - different variables",
			input:    lines("\tlda x", "\tsta y"),
			expected: 2, // different
		},
		{
			name:     "not self - immediate load",
			input:    lines("\tlda #$05", "\tsta #$05"),
			expected: 2, // immediate + immediate doesn't match
		},
		{
			name:     "not self - SM code pattern (+ offset)",
			input:    lines("\tlda x", "\tsta _L1+1"),
			expected: 2, // SM code pattern
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			parsed := parseLines(tt.input)
			result := passSelfAssignment(parsed)
			cleaned := stripOptMarkers(result)
			if got := len(cleaned); got != tt.expected {
				t.Errorf("got %d lines, want %d\ninput: %v\noutput: %v", got, tt.expected, tt.input, linesToString(cleaned))
			}
		})
	}
}

func TestPassLoadNoCommentReset(t *testing.T) {
	// Comments should NOT reset register state
	parsed := parseLines(lines(
		"\tlda b",
		"\tsta a",
		"; comment",
		"\tlda b", // should be redundant
	))
	result := passLoadElimination(parsed)
	if len(result) != 3 {
		t.Errorf("expected 3 lines (comment kept, lda b removed), got %d:\n%v", len(result), linesToString(result))
	}
}

func TestOptimizeIntegration(t *testing.T) {
	input := lines(
		"; a = b",
		"; @@OPT:LINEAR:LET",
		"\tlda b",
		"\tsta a",
		"; second use of b (redundant load)",
		"; @@OPT:LINEAR:LET",
		"\tlda b",
		"\tsta c",
	)

	cfg := &Config{EnableLoad: true}
	output := Optimize(input, cfg)

	// 2 source comments + 3 asm lines (lda b removed) = 5
	if len(output) != 5 {
		t.Errorf("expected 5 lines, got %d:\n%v", len(output), output)
	}
}

func TestOptimizeWithDebug(t *testing.T) {
	input := lines(
		"\tlda x",
		"; @@OPT:LINEAR:LET",
		"\tlda x",
		"\tsta y",
	)

	cfg := &Config{EnableLoad: true, Debug: true}
	output := Optimize(input, cfg)

	// Header + 2 kept lines + 1 removed annotation = 4
	if len(output) != 4 {
		t.Errorf("expected 4 lines, got %d:\n%v", len(output), output)
	}
	if !containsPrefix(output, "; --- peephole:") {
		t.Errorf("expected peephole summary line:\n%v", output)
	}
	if !containsSubstring(output, "[removed]") {
		t.Errorf("expected [removed] annotation:\n%v", output)
	}
}

func containsPrefix(lines []string, prefix string) bool {
	for _, l := range lines {
		if len(l) >= len(prefix) && l[:len(prefix)] == prefix {
			return true
		}
	}
	return false
}

func containsSubstring(lines []string, sub string) bool {
	for _, l := range lines {
		for i := 0; i <= len(l)-len(sub); i++ {
			if l[i:i+len(sub)] == sub {
				return true
			}
		}
	}
	return false
}

func TestPassStoreReload(t *testing.T) {
	cfg := &Config{}

	tests := []struct {
		name     string
		input    []string
		expected int
	}{
		{
			name:     "sta x then lda x",
			input:    lines("\tsta x", "\tlda x"),
			expected: 1,
		},
		{
			name:     "different variables",
			input:    lines("\tsta x", "\tlda y"),
			expected: 2,
		},
		{
			name:     "ldy between sta and lda",
			input:    lines("\tsta x", "\tldy #0", "\tlda x"),
			expected: 2,
		},
		{
			name:     "adc between — modifies A",
			input:    lines("\tsta x", "\tadc #1", "\tlda x"),
			expected: 3,
		},
		{
			name:     "label between — barrier",
			input:    lines("\tsta x", "label", "\tlda x"),
			expected: 3,
		},
		{
			name:     "different registers",
			input:    lines("\tsta x", "\tldx x"),
			expected: 2,
		},
		{
			name:     "comment between",
			input:    lines("\tsta x", "; source line", "\tlda x"),
			expected: 2,
		},
		{
			name:     "comment and ldy between",
			input:    lines("\tsta x", "; POKE addr", "\tldy #5", "\tlda x"),
			expected: 3,
		},
		{
			name:     "SM code — operand with +",
			input:    lines("\tsta _L1+1", "\tlda _L1+1"),
			expected: 2,
		},
		{
			name:     "indexed store — operand with (",
			input:    lines("\tsta (zp),y", "\tlda (zp),y"),
			expected: 2,
		},
		{
			name:     "indexed store — operand with ,",
			input:    lines("\tsta $D020,x", "\tlda $D020,x"),
			expected: 2,
		},
		{
			name:     "different operand offsets",
			input:    lines("\tsta x", "\tlda x+1"),
			expected: 2,
		},
		{
			name:     "inx between — does not modify A",
			input:    lines("\tsta x", "\tinx", "\tlda x"),
			expected: 2,
		},
		{
			name:     "dex between — does not modify A",
			input:    lines("\tsta x", "\tdex", "\tlda x"),
			expected: 2,
		},
		{
			name:     "iny between — does not modify A",
			input:    lines("\tsta x", "\tiny", "\tlda x"),
			expected: 2,
		},
		{
			name:     "dey between — does not modify A",
			input:    lines("\tsta x", "\tdey", "\tlda x"),
			expected: 2,
		},
		{
			name:     "tax between — does not modify A",
			input:    lines("\tsta x", "\ttax", "\tlda x"),
			expected: 2,
		},
		{
			name:     "tay between — does not modify A",
			input:    lines("\tsta x", "\ttay", "\tlda x"),
			expected: 2,
		},
		{
			name:     "txa between — modifies A, keep lda",
			input:    lines("\tsta x", "\ttxa", "\tlda x"),
			expected: 3,
		},
		{
			name:     "tya between — modifies A, keep lda",
			input:    lines("\tsta x", "\ttya", "\tlda x"),
			expected: 3,
		},
		{
			name:     "pla between — modifies A, keep lda",
			input:    lines("\tsta x", "\tpla", "\tlda x"),
			expected: 3,
		},
		{
			name:     "multiple stores — pick the last sta",
			input:    lines("\tsta x", "\tsta y", "\tlda x"),
			expected: 2, // sta x kept, sta y kept, lda x removed
		},
		{
			name:     "@@OPT marker between — skip, lda removed",
			input:    lines("\tsta x", "; @@OPT:LINEAR:LET", "\tlda x"),
			expected: 1, // @@OPT stripped by stripOptMarkers
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			parsed := parseLines(tt.input)
			result := passStoreReload(parsed, cfg)
			cleaned := stripOptMarkers(result)
			if got := len(cleaned); got != tt.expected {
				t.Errorf("got %d lines, want %d\ninput: %v\noutput: %v", got, tt.expected, tt.input, linesToString(cleaned))
			}
		})
	}
}

func TestPassStoreReloadIO(t *testing.T) {
	cfg := &Config{}
	cfg.IOMap[0xD020] = true

	parsed := parseLines(lines("\tsta $D020", "\tlda $D020"))
	result := passStoreReload(parsed, cfg)
	cleaned := stripOptMarkers(result)
	if len(cleaned) != 2 {
		t.Errorf("expected 2 lines (I/O skip), got %d", len(cleaned))
	}
}