package container

import (
	"iter"
	"sync"
)

// TransformSlice returns a new slice containing the results of applying the function f to each item in the original slice.
func TransformSlice[T any, O any](s []T, f func(e T) O) []O {
	r := make([]O, len(s))
	for i, e := range s {
		r[i] = f(e)
	}
	return r
}

// SliceIterTransform is TransformSlice with iter.Seq
func SliceIterTransform[T any, O any](s iter.Seq[T], f func(e T) O) iter.Seq[O] {
	return func(yield func(O) bool) {
		for e := range s {
			if !yield(f(e)) {
				return
			}
		}
	}
}

// SliceIterFilter return an iterator of the slice filtered by the function filter.
func SliceIterFilter[T any](s iter.Seq[T], filter func(e T) bool) iter.Seq[T] {
	return func(yield func(T) bool) {
		for e := range s {
			if filter(e) {
				continue
			}
			if !yield(e) {
				return
			}
		}
	}
}

// SliceToMap converts a slice to a map base on the result of the function f.
func SliceToMap[T any, K comparable](s []T, f func(e T) K) map[K]T {
	m := make(map[K]T, len(s))
	for _, e := range s {
		m[f(e)] = e
	}

	return m
}

// ForEach is a shortcut for for...range.
func ForEach[T any](s []T, f func(i int, e T)) {
	for i, e := range s {
		f(i, e)
	}
}

// Slice is a thread-safe generic slice.
type Slice[T any] struct {
	s  []T
	mu *sync.Mutex
}

func NewSlice[T any]() *Slice[T] {
	return &Slice[T]{
		s:  make([]T, 0),
		mu: &sync.Mutex{},
	}
}

func (s *Slice[T]) Append(e ...T) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.s = append(s.s, e...)
}

func (s *Slice[T]) Len() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	return len(s.s)
}

func (s *Slice[T]) Get(i int) T {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.s[i]
}

func (s *Slice[T]) Dump() []T {
	s.mu.Lock()
	defer s.mu.Unlock()
	r := make([]T, len(s.s))
	copy(r, s.s)
	return r
}

// All return a iterator of the slice and the code in for...range is in lock.
//
// For example:
//
//	for i, e := range s.All() { // do something in lock }
func (s *Slice[T]) All() iter.Seq2[int, T] {
	return func(yield func(int, T) bool) {
		s.mu.Lock()
		defer s.mu.Unlock()

		for i, v := range s.s {
			if !yield(i, v) {
				return
			}
		}
	}
}

// Values is similar to All and returns keys of the slice.
//
// For example:
//
//	for e := range s.Values() { // do something in lock }
func (s *Slice[T]) Values() iter.Seq[T] {
	return func(yield func(T) bool) {
		s.mu.Lock()
		defer s.mu.Unlock()

		for _, v := range s.s {
			if !yield(v) {
				return
			}
		}
	}
}