package testutil

import (
	"bytes"
	"context"
	"strings"
	"sync"
	"testing"
)

// WaitBuffer is a thread-safe buffer (io.Writer) that supports
// blocking until the accumulated content matches a condition.
// It is intended for tests that need to wait for specific output
// from a command or process before proceeding.
//
// WaitBuffer is safe for concurrent use. Multiple goroutines may
// write to it, and WaitFor/WaitForCond may be called from any
// goroutine.
type WaitBuffer struct {
	mu      sync.Mutex
	buf     bytes.Buffer
	waiters []*wbWaiter
}

type wbWaiter struct {
	cond func(string) bool
	ch   chan struct{}
	once sync.Once
}

// NewWaitBuffer returns a new WaitBuffer. It can be used as a
// plain thread-safe io.Writer even if WaitFor is never called.
func NewWaitBuffer() *WaitBuffer {
	return &WaitBuffer{}
}

// Write implements io.Writer. It is safe for concurrent use.
func (wb *WaitBuffer) Write(p []byte) (int, error) {
	wb.mu.Lock()
	defer wb.mu.Unlock()

	n, err := wb.buf.Write(p)
	s := wb.buf.String()
	for _, w := range wb.waiters {
		if w.cond(s) {
			w.once.Do(func() { close(w.ch) })
		}
	}
	return n, err
}

// WaitFor blocks until the accumulated output contains signal or
// ctx expires. Returns nil on match, ctx.Err() on timeout.
// Safe to call from any goroutine.
func (wb *WaitBuffer) WaitFor(ctx context.Context, signal string) error {
	return wb.WaitForNth(ctx, signal, 1)
}

// WaitForNth blocks until the accumulated output contains at least
// n occurrences of signal, or ctx expires. Returns nil on match,
// ctx.Err() on timeout. Safe to call from any goroutine.
func (wb *WaitBuffer) WaitForNth(ctx context.Context, signal string, n int) error {
	return wb.WaitForCond(ctx, func(s string) bool {
		return strings.Count(s, signal) >= n
	})
}

// WaitForCond blocks until cond returns true for the accumulated
// output, or ctx expires. Returns nil on match, ctx.Err() on
// timeout. Safe to call from any goroutine.
func (wb *WaitBuffer) WaitForCond(ctx context.Context, cond func(string) bool) error {
	wb.mu.Lock()
	if cond(wb.buf.String()) {
		wb.mu.Unlock()
		return nil
	}
	w := &wbWaiter{
		cond: cond,
		ch:   make(chan struct{}),
	}
	wb.waiters = append(wb.waiters, w)
	wb.mu.Unlock()

	select {
	case <-w.ch:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}

// RequireWaitFor blocks until the accumulated output contains
// signal or ctx expires. On timeout, fails the test with a
// message showing what was expected and what was written so far.
//
// Safety: Must only be called from the Go routine that created
// `t`.
func (wb *WaitBuffer) RequireWaitFor(ctx context.Context, t testing.TB, signal string) {
	t.Helper()
	if err := wb.WaitFor(ctx, signal); err != nil {
		t.Fatalf("WaitBuffer: signal %q not found; buffer contents:\n%s", signal, wb.String())
	}
}

// Bytes returns a copy of the accumulated output.
func (wb *WaitBuffer) Bytes() []byte {
	wb.mu.Lock()
	defer wb.mu.Unlock()
	return bytes.Clone(wb.buf.Bytes())
}

// String returns the accumulated output as a string.
func (wb *WaitBuffer) String() string {
	wb.mu.Lock()
	defer wb.mu.Unlock()
	return wb.buf.String()
}