package itertools

import (
	"iter"

	"golang.org/x/exp/constraints"
)

// Number is the set of types Count and CountStep can generate, i.e. anything
// that supports the + operator with an integer or floating-point step.
type Number interface {
	constraints.Integer | constraints.Float
}

// Take returns a lazy iterator over the first n items of seq.
// It stops seq early once n items have been yielded,
// so it is safe to apply to infinite sequences like Count.
// Use slices.Collect to build a slice with the results.
func Take[T any](n int, seq iter.Seq[T]) iter.Seq[T] {
	return func(yield func(T) bool) {
		taken := 0
		for v := range seq {
			if taken >= n {
				return
			}
			if !yield(v) {
				return
			}
			taken++
		}
	}
}

// RepeatForever returns an infinite iterator that yields v forever.
// Use Repeat for the bounded form.
// The caller must break or use a bounded helper (e.g. Take) to stop it.
func RepeatForever[T any](v T) iter.Seq[T] {
	return func(yield func(T) bool) {
		for yield(v) {
		}
	}
}

// Repeat returns an iterator that yields v exactly times times.
// A non-positive times yields nothing.
func Repeat[T any](v T, times int) iter.Seq[T] {
	return func(yield func(T) bool) {
		for range times {
			if !yield(v) {
				return
			}
		}
	}
}

// Count returns an infinite iterator that yields start, start+1, start+2, …
// The caller must break or use a bounded helper (e.g. Take) to stop it.
func Count[N Number](start N) iter.Seq[N] {
	return CountStep(start, 1)
}

// CountStep returns an infinite iterator that yields start, start+step,
// start+2*step, …
// The caller must break or use a bounded helper (e.g. Take) to stop it.
func CountStep[N Number](start, step N) iter.Seq[N] {
	return func(yield func(N) bool) {
		for i := start; ; i += step {
			if !yield(i) {
				return
			}
		}
	}
}

// Cycle returns an infinite iterator that yields the items of seq forever.
// Iterates over seq once, saving a copy of each item to yield each later.
// Memory grows with the length of seq; an infinite seq makes the saved slice
// grow without bound and never reaches the replay loop.
// An empty seq yields nothing.
// The caller must break or use a bounded helper (e.g. Take) to stop it.
func Cycle[T any](seq iter.Seq[T]) iter.Seq[T] {
	return func(yield func(T) bool) {
		var saved []T
		for v := range seq {
			if !yield(v) {
				return
			}
			saved = append(saved, v)
		}
		if len(saved) == 0 {
			return
		}
		for {
			for _, v := range saved {
				if !yield(v) {
					return
				}
			}
		}
	}
}

// Chain returns an iterator that yields the items of each seq in turn,
// exhausting the first before moving to the second, and so on.
// With no arguments it yields nothing.
func Chain[T any](seqs ...iter.Seq[T]) iter.Seq[T] {
	return func(yield func(T) bool) {
		for _, seq := range seqs {
			for v := range seq {
				if !yield(v) {
					return
				}
			}
		}
	}
}

// Filter returns an iterator that yields only the items of seq for which
// predicate returns true.
func Filter[T any](predicate func(T) bool, seq iter.Seq[T]) iter.Seq[T] {
	return func(yield func(T) bool) {
		for v := range seq {
			if predicate(v) {
				if !yield(v) {
					return
				}
			}
		}
	}
}

// DropWhile discards items from seq while the predicate returns true.
// Then it yields all remaining items, regardless of the predicate.
func DropWhile[T any](predicate func(T) bool, seq iter.Seq[T]) iter.Seq[T] {
	return func(yield func(T) bool) {
		dropping := true
		for v := range seq {
			if dropping {
				if predicate(v) {
					continue
				}
				dropping = false
			}
			if !yield(v) {
				return
			}
		}
	}
}