package llm

import (
	"strings"
	"testing"
)

func run(chunks ...string) (string, string) {
	var f thinkFilter
	var streamed strings.Builder
	for _, c := range chunks {
		streamed.WriteString(f.push(c))
	}
	streamed.WriteString(f.finish())
	return streamed.String(), f.out.String()
}

func TestThinkFilter(t *testing.T) {
	cases := []struct {
		name   string
		chunks []string
		want   string
	}{
		{"plain", []string{"Hallo ", "Welt"}, "Hallo Welt"},
		{"block", []string{"<think>reasoning\n</think>\n\nHallo!"}, "Hallo!"},
		{"split tags", []string{"<th", "ink>abc</th", "ink>", "\n\nHal", "lo"}, "Hallo"},
		{"char by char", strings.Split("<think>x</think>Ja", ""), "Ja"},
		{"stray close", []string{"Hallo!</think>\n\nHallo!"}, "Hallo!\n\nHallo!"},
		{"unclosed with gap (Dahl)", []string{"<think>The user asks.\n\n\nHallo! Wie geht's?"}, "Hallo! Wie geht's?"},
		{"gap streams live", []string{"<think>a\n", "\n", "\nAntw", "ort"}, "Antwort"},
		{"gap outside block kept", []string{"A\n\n\nB"}, "A\n\n\nB"},
		{"unclosed no gap", []string{"<think>nur Gedanken"}, "nur Gedanken"},
		{"lt not a tag", []string{"a <b> c <", "3"}, "a <b> c <3"},
		{"trailing partial", []string{"x <thi"}, "x <thi"},
	}
	for _, c := range cases {
		streamed, out := run(c.chunks...)
		if streamed != c.want || out != c.want {
			t.Errorf("%s: streamed=%q out=%q want %q", c.name, streamed, out, c.want)
		}
	}
}

func TestParseStreamStripsThink(t *testing.T) {
	sse := `data: {"choices":[{"delta":{"content":"","role":"assistant"}}]}

data: {"choices":[{"delta":{"reasoning":"hidden"}}]}

data: {"choices":[{"delta":{"content":"<think>The user"}}]}

data: {"choices":[{"delta":{"content":" greets.\n</thi"}}]}

data: {"choices":[{"delta":{"content":"nk>\n\nHallo!"}}]}

data: {"choices":[{"delta":{},"finish_reason":"stop"}]}

data: [DONE]
`
	var deltas []string
	resp, err := parseStream(strings.NewReader(sse), func(s string) { deltas = append(deltas, s) })
	if err != nil || resp.Content != "Hallo!" || strings.Join(deltas, "") != "Hallo!" {
		t.Fatalf("%q %v %v", resp.Content, deltas, err)
	}
}

// The answer after a swallowed </think> must be streamed incrementally, not only at the end.
func TestThinkGapStreamsIncrementally(t *testing.T) {
	var f thinkFilter
	if v := f.push("<think>Überlegung\n\n\nHal"); v != "Hal" {
		t.Fatalf("got %q", v)
	}
	if v := f.push("lo"); v != "lo" {
		t.Fatalf("got %q", v)
	}
}