package main

import "testing"

func TestStripCompletionNoise(t *testing.T) {
	tests := []struct {
		name string
		in   string
		out  string
	}{
		{"plain", "hello\nworld", "hello\nworld"},
		{"markdown fences", "```go\nfmt.Println()\n```", "fmt.Println()"},
		{"leaked prefix tag", "<prefix>\ncode\n</prefix>", "code"},
		{"leaked suffix tag", "<suffix>\ncode\n</suffix>", "code"},
		{"info lines", ":: cost=0.01\nresult", "result"},
		{"mixed", "```\n:: info\nreal code\n<prefix>\n```", "real code"},
		{"empty", "", ""},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := stripCompletionNoise(tt.in)
			if got != tt.out {
				t.Errorf("stripCompletionNoise(%q) = %q, want %q", tt.in, got, tt.out)
			}
		})
	}
}

func TestStripPrefixEcho(t *testing.T) {
	tests := []struct {
		name   string
		prefix string
		result string
		out    string
	}{
		{"no echo", "func main() {", "\n\tfmt.Println()", "\n\tfmt.Println()"},
		{"full echo", "func main() {", "func main() {\n\tfmt.Println()", "\n\tfmt.Println()"},
		{"partial tail echo", "hello world", "world new stuff", " new stuff"},
		{"empty prefix", "", "some result", "some result"},
		{"empty result", "prefix", "", ""},
		{"long prefix short echo", "a]" + string(make([]byte, 300)), "result", "result"},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := stripPrefixEcho(tt.prefix, tt.result)
			if got != tt.out {
				t.Errorf("stripPrefixEcho(%q, %q) = %q, want %q", tt.prefix, tt.result, got, tt.out)
			}
		})
	}
}