package render

import (
	"context"
	"io"
	"net/http"
	"testing"

	"github.com/a-h/templ"
)

type writeCountingResponseWriter struct {
	header http.Header
	writes int
	body   []byte
}

func (w *writeCountingResponseWriter) Header() http.Header {
	if w.header == nil {
		w.header = make(http.Header)
	}
	return w.header
}

func (w *writeCountingResponseWriter) Write(p []byte) (int, error) {
	w.writes++
	w.body = append(w.body, p...)
	return len(p), nil
}

func (w *writeCountingResponseWriter) WriteHeader(statusCode int) {}

func (w *writeCountingResponseWriter) Flush() {}

func TestStreamComponentWritesIncrementally(t *testing.T) {
	w := &writeCountingResponseWriter{}
	component := templ.ComponentFunc(func(ctx context.Context, writer io.Writer) error {
		if _, err := writer.Write([]byte("a")); err != nil {
			return err
		}
		if _, err := writer.Write([]byte("b")); err != nil {
			return err
		}
		return nil
	})

	if err := streamComponent(context.Background(), w, component); err != nil {
		t.Fatalf("streamComponent() error = %v", err)
	}
	if got := string(w.body); got != "ab" {
		t.Fatalf("body = %q, want %q", got, "ab")
	}
	if w.writes < 2 {
		t.Fatalf("writes = %d, want at least 2 incremental writes", w.writes)
	}
}