package iter

// ...

// ChunkIter works like slices.Chunk but with iterators,
// i. e. iterates over it and yields batches of size n.
// Panics if n is less than 1.
func ChunkIter[I Seq[V], V any](it I, n int) Seq[[]V] {
	if n < 1 {
		panic("cannot be less than 1")
	}

	return func(yield func([]V) bool) {
		batch := []V{}
		for v := range it {
			batch = append(batch, v)
			if len(batch) == n {
				if !yield(batch) {
					return
				}

				batch = []V{}
			}
		}

		if len(batch) > 0 {
			yield(batch)
		}
	}
}

// ChunkIter2 works like slices.ChunkIter,
// but yields pairs of batches of size n.
// Panics if n is less than 1.
func ChunkIter2[I Seq2[K, V], K any, V any](it I, n int) Seq2[[]K, []V] {
	if n < 1 {
		panic("cannot be less than 1")
	}

	return func(yield func([]K, []V) bool) {
		kbatch := []K{}
		vbatch := []V{}

		for k, v := range it {
			kbatch = append(kbatch, k)
			vbatch = append(vbatch, v)

			if len(kbatch) == n {
				if !yield(kbatch, vbatch) {
					return
				}

				kbatch = []K{}
				vbatch = []V{}
			}
		}

		if len(kbatch) > 0 {
			yield(kbatch, vbatch)
		}
	}
}