package stable

import (
	"iter"
)

const DefaultElementsPerBucket = 32

// Array is a resizable array whose values will never move in memory.
// This means it is safe to take a pointer to a value within the array
// while continuing to append to it.
type Array[T any] struct {
	buckets             []bucket[T]
	last                int
	elements_per_bucket int
}

func (s *Array[T]) Init() {
	s.InitWithCapacity(DefaultElementsPerBucket)
}

func (s *Array[T]) InitWithCapacity(elements_per_bucket int) {
	if elements_per_bucket <= 0 {
		elements_per_bucket = DefaultElementsPerBucket
	}

	s.elements_per_bucket = elements_per_bucket

	s.buckets = s.buckets[:0]
	s.buckets = append(s.buckets, make(bucket[T], 0, s.elements_per_bucket))
	s.last = 0
}

func (s *Array[T]) Reset() {
	s.buckets = s.buckets[:0]
	s.last = 0
}

func (s *Array[T]) Append(value T) *T {
	if len(s.buckets) == 0 {
		s.Init()
	}

	if len(s.buckets[s.last]) == cap(s.buckets[s.last]) {
		s.buckets = append(s.buckets, make(bucket[T], 0, s.elements_per_bucket))
		s.last += 1
	}

	s.buckets[s.last] = append(s.buckets[s.last], value)
	return &s.buckets[s.last][len(s.buckets[s.last])-1]
}

func (s *Array[T]) AppendMany(values ...T) (first *T) {
	if len(values) == 0 {
		return nil
	}

	first = s.Append(values[0])

	if len(values) > 1 {
		for _, v := range values[1:] {
			s.Append(v)
		}
	}

	return
}

func (s *Array[T]) Get(index int) *T {
	b := s.buckets[index/s.elements_per_bucket]
	return &b[index%s.elements_per_bucket]
}

func (s *Array[T]) Set(index int, value T) {
	*s.Get(index) = value
}

func (s *Array[T]) Len() int {
	return s.Cap() - (cap(s.buckets[s.last]) - len(s.buckets[s.last]))
}

func (s *Array[T]) Cap() int {
	return len(s.buckets) * s.elements_per_bucket
}

func (s *Array[T]) Pointers() iter.Seq2[int, *T] {
	return func(yield func(int, *T) bool) {
		for bi := range s.buckets {
			startIdx := bi * s.elements_per_bucket
			for i := range s.buckets[bi] {
				if !yield(startIdx+i, &s.buckets[bi][i]) {
					return
				}
			}
		}
	}
}

func (s *Array[T]) Values() iter.Seq2[int, T] {
	return func(yield func(int, T) bool) {
		for bi, b := range s.buckets {
			startIdx := bi * s.elements_per_bucket
			for i := range b {
				if !yield(startIdx+i, b[i]) {
					return
				}
			}
		}
	}
}

type bucket[T any] = []T