package stable

import (
	"iter"
	"math/bits"
	"unsafe"

	"git.brut.systems/judah/xx"
)

const (
	DefaultChunkSizeShift = 4
)

// Xar is an implementation of Andrew Reece's Xar (Exponential Array).
//
// See: https://www.youtube.com/watch?v=i-h95QIGchY
type Xar[T any] struct {
	shift  uint8
	count  uint64
	chunks [][]T
}

func (x *Xar[T]) Init() {
	x.InitWithSize(DefaultChunkSizeShift)
}

func (x *Xar[T]) InitWithSize(size_shift uint8) {
	if len(x.chunks) != 0 {
		x.Reset()
	}

	x.shift = size_shift
}

func (x *Xar[T]) Reset() {
	for _, c := range x.chunks {
		clear(c)
	}

	x.count = 0
}

func (x *Xar[T]) Append(value T) *T {
	if x.shift == 0 {
		x.Init()
	}

	chunk_idx, idx_in_chunk, chunk_cap := x.getChunk(x.count)
	x.count += 1
	if chunk_idx >= uint64(len(x.chunks)) {
		x.chunks = append(x.chunks, make([]T, chunk_cap))
	}

	slot := &x.chunks[chunk_idx][idx_in_chunk]
	*slot = value
	return slot
}

func (x *Xar[T]) AppendMany(values ...T) *T {
	if len(values) == 0 {
		return nil
	}

	first := x.Append(values[0])

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

	return first
}

func (x *Xar[T]) Get(index int) *T {
	chunk_idx, idx_in_chunk, _ := x.getChunk(uint64(index))
	return &x.chunks[chunk_idx][idx_in_chunk]
}

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

func (x *Xar[T]) Remove(index int) {
	x.Set(index, *x.Get(int(x.count - 1)))
	x.count -= 1
}

func (x *Xar[T]) Len() int {
	return int(x.count)
}

func (x *Xar[T]) Cap() (l int) {
	for _, c := range x.chunks {
		l += cap(c)
	}
	return
}

func (x *Xar[T]) Pointers() iter.Seq2[int, *T] {
	return func(yield func(int, *T) bool) {
		idx := -1
		for chunk_idx, idx_in_chunk := range x.iter() {
			idx += 1
			if !yield(idx, &x.chunks[chunk_idx][idx_in_chunk]) {
				break
			}
		}
	}
}

func (x *Xar[T]) Values() iter.Seq2[int, T] {
	return func(yield func(int, T) bool) {
		idx := -1
		for chunk_idx, idx_in_chunk := range x.iter() {
			idx += 1
			if !yield(idx, x.chunks[chunk_idx][idx_in_chunk]) {
				break
			}
		}
	}
}

func (x *Xar[T]) iter() iter.Seq2[uint64, uint64] {
	return func(yield func(uint64, uint64) bool) {
		chunk_size := 1 << x.shift
	outer:
		for chunk_idx := range x.chunks {
			for idx_in_chunk := range chunk_size - 1 {
				if uint64(chunk_idx+idx_in_chunk) >= uint64(x.count) {
					break outer
				}

				if !yield(uint64(chunk_idx), uint64(idx_in_chunk)) {
					break outer
				}
			}

			chunk_size <<= xx.BoolUint(chunk_idx > 0)
		}
	}
}

func (x *Xar[T]) getChunk(index uint64) (chunk_idx uint64, idx_in_chunk uint64, chunk_cap uint64) {
	if true /* branchless */ {
		var (
			i_shift  = index >> x.shift
			i_shift2 = i_shift != 0
			msb      = msb64(i_shift | 1)
			b        = uint64(*(*uint8)(unsafe.Pointer(&i_shift2)))
		)

		chunk_idx = msb + b
		idx_in_chunk = index - (b << (msb + uint64(x.shift)))
		chunk_cap = 1 << (msb + uint64(x.shift))
	} else {
		idx_in_chunk = index
		chunk_cap = 1 << x.shift

		i_shift := index >> x.shift
		if i_shift > 0 {
			chunk_idx = msb64(i_shift | 1)
			chunk_cap = 1 << (chunk_idx + uint64(x.shift))
			idx_in_chunk -= chunk_cap
			chunk_idx += 1
		}
	}

	return
}

func msb64(x uint64) uint64 {
	return uint64(63 - bits.LeadingZeros64(x))
}