package pipeline

import (
	"context"
	"time"
)

// Batch collects items into slices of up to size, flushing when full or when timeout elapses.
func Batch[T any](ctx context.Context, in <-chan T, size int, timeout time.Duration) <-chan []T {
	if in == nil {
		panic("pipeline: Batch input channel must not be nil")
	}
	if size < 1 {
		panic("pipeline: Batch size must be at least 1")
	}

	out := make(chan []T)

	go func() {
		defer close(out)

		buf := make([]T, 0, size)
		ticker := time.NewTicker(timeout)
		defer ticker.Stop()

		flush := func() bool {
			if len(buf) == 0 {
				return true
			}
			batch := buf
			buf = make([]T, 0, size)
			select {
			case out <- batch:
				return true
			case <-ctx.Done():
				return false
			}
		}

		for {
			select {
			case v, ok := <-in:
				if !ok {
					flush()
					return
				}
				buf = append(buf, v)
				if len(buf) >= size {
					if !flush() {
						return
					}
				}
			case <-ticker.C:
				if !flush() {
					return
				}
			case <-ctx.Done():
				return
			}
		}
	}()

	return out
}