package xarr

import "base:builtin"
import "base:intrinsics"
import "common:relptr"
import "core:mem"

BASE_CHUNK_SIZE :: uint(64)
BASE_CHUNK_SIZE_LOG2 :: intrinsics.constant_log2(BASE_CHUNK_SIZE)
BASE_CHUNK_SHIFT :: BASE_CHUNK_SIZE_LOG2 - 1
NUM_CHUNKS :: 30 when (size_of(uint) == 8) else 26 // on 32 bit systems max size is 0x80000000 which is about half the addressable space

Xarr :: struct($T: typeid, $SOA := false) {
	len:                   int,
	allocated_chunks_mask: u32,
	chunks:                ([NUM_CHUNKS]relptr.SOA_Slice(T) when SOA else [NUM_CHUNKS]relptr.Ptr(
			T,
		)),
}

UINT_BITS :: size_of(uint) * 8

msb :: #force_inline proc "contextless" (#any_int idx: uint) -> i32 {
	return i32(UINT_BITS - intrinsics.count_leading_zeros(idx)) - 1
}

chunk_by_index :: #force_inline proc "contextless" (#any_int idx: uint) -> (chunk: i32) {
	return max(msb(idx) - BASE_CHUNK_SHIFT, 0)
}

chunk_size :: #force_inline proc "contextless" (chunk_idx: i32) -> uint {
	return BASE_CHUNK_SIZE << intrinsics.saturating_sub(u32(chunk_idx), 1)
}

get_chunk_slice_scalar :: #force_inline proc(
	a: $T/Xarr($E, false),
	chunk_idx: i32,
	base := context.user_ptr,
) -> []E {
	return relptr.deref_multi_ptr(a.chunks[chunk_idx], base)[:chunk_size(chunk_idx)]
}

get_chunk_slice_soa :: #force_inline proc(
	a: $T/Xarr($E, true),
	chunk_idx: i32,
	base := context.user_ptr,
) -> #soa[]E {
	return relptr.deref_soa_slice(a.chunks[chunk_idx], base)
}

get_chunk_slice :: proc {
	get_chunk_slice_scalar,
	get_chunk_slice_soa,
}

capacity_from_allocated_mask :: #force_inline proc "contextless" (allocated_mask: u32) -> uint {
	return(
		uint(allocated_mask >> 1) << BASE_CHUNK_SIZE_LOG2 +
		uint(allocated_mask & 1) << BASE_CHUNK_SIZE_LOG2 \
	)
}

capacity :: #force_inline proc "contextless" (a: $T/Xarr($E, $SOA)) -> uint {
	allocated_mask := a.allocated_chunks_mask
	return capacity_from_allocated_mask(allocated_mask)
}

len :: #force_inline proc "contextless" (a: $T/Xarr($E, $SOA)) -> int {
	return a.len
}

reserve :: proc(
	a: $T/^Xarr($E, $SOA),
	cap: int,
	allocator := context.allocator,
	base := context.user_ptr,
) #no_bounds_check {
	allocated_mask := a.allocated_chunks_mask

	current_chunk := msb(allocated_mask)
	required_chunks := chunk_by_index(max(cap - 1, 0)) + 1

	assert(required_chunks <= NUM_CHUNKS)

	for i := current_chunk + 1; i < required_chunks; i += 1 {
		when SOA {
			chunk_slice := make_soa_slice(#soa[]E, chunk_size(i), allocator)
			a.chunks[i] = relptr.from_soa_slice(chunk_slice, base)
		} else {
			chunk_slice := make([]E, chunk_size(i), allocator)
			a.chunks[i] = relptr.from_multi_ptr(raw_data(chunk_slice), base)
		}
		a.allocated_chunks_mask |= u32(1) << u8(i)
	}
}

resize :: proc(
	a: $T/^Xarr($E, $SOA),
	new_len: int,
	allocator := context.allocator,
	base := context.user_ptr,
) {
	reserve(a, new_len, allocator, base)
	a.len = new_len
}

append_elem :: proc(
	a: $T/^Xarr($E, $SOA),
	elem: E,
	allocator := context.allocator,
	base := context.user_ptr,
) {
	if capacity(a^) < uint(a.len + 1) {
		reserve(a, a.len + 1, allocator, base)
	}
	#no_bounds_check {
		chunk_idx, idx_within_chunk := translate_index(a.len)
		when SOA {
			slice := relptr.deref_soa_slice(a.chunks[chunk_idx], base)
			slice[idx_within_chunk] = elem
		} else {
			relptr.deref_multi_ptr(a.chunks[chunk_idx], base)[idx_within_chunk] = elem
		}
	}
	a.len += 1
}

append_elems :: proc(
	a: $T/^Xarr($E, $SOA),
	elems: ..E,
	allocator := context.allocator,
	base := context.user_ptr,
) {
	if builtin.len(elems) == 0 {
		return
	}

	if capacity(a^) < uint(a.len + builtin.len(elems)) {
		reserve(a, a.len + builtin.len(elems), allocator, base)
	}
	set_elems_assume_allocated(a, elems, base)
	a.len += builtin.len(elems)
}

append :: proc {
	append_elem,
	append_elems,
}

translate_index :: #force_inline proc "contextless" (
	#any_int idx: int,
) -> (
	chunk_idx: i32,
	idx_within_chunk: uint,
) {
	chunk_idx = chunk_by_index(idx)
	idx_within_chunk = uint(idx) & (chunk_size(chunk_idx) - 1)

	return
}

@(private = "file")
set_elems_assume_allocated :: proc(
	a: $T/^Xarr($E, $SOA),
	elems: []E,
	base: rawptr,
) #no_bounds_check {
	for &e, i in elems {
		idx := a.len + i
		chunk_idx, idx_within_chunk := translate_index(idx)

		when SOA {
			slice := relptr.deref_soa_slice(a.chunks[chunk_idx], base)
			slice[idx_within_chunk] = e
		} else {
			relptr.deref_multi_ptr(a.chunks[chunk_idx], base)[idx_within_chunk] = e
		}
	}
}

set :: proc(a: $T/Xarr($E, $SOA), #any_int idx: int, val: E) {
	assert(idx >= 0 && idx < a.len)
	chunk_idx, idx_within_chunk := translate_index(idx)
	return get_chunk_slice(a, chunk_idx)[idx_within_chunk]
}

get :: proc(a: $T/Xarr($E, $SOA), #any_int idx: int) -> E {
	assert(idx >= 0 && idx < a.len)

	chunk_idx, idx_within_chunk := translate_index(idx)
	return get_chunk_slice(a, chunk_idx)[idx_within_chunk]
}

get_ptr_scalar :: proc(a: $T/^Xarr($E, false), #any_int idx: int) -> ^E {
	assert(idx >= 0 && idx < a.len)

	chunk_idx, idx_within_chunk := translate_index(idx)
	return &get_chunk_slice_scalar(a, chunk_idx)[idx_within_chunk]
}

get_ptr_soa :: proc(a: $T/^Xarr($E, true), #any_int idx: int) -> #soa^#soa[]E {
	assert(idx >= 0 && idx < a.len)

	chunk_idx, idx_within_chunk := translate_index(idx)
	return &get_chunk_slice_soa(a, chunk_idx)[idx_within_chunk]
}

get_ptr :: proc {
	get_ptr_scalar,
	get_ptr_soa,
}

unordered_remove :: proc(a: $T/^Xarr($E, $SOA), #any_int idx: int) {
	assert(idx >= 0 && idx < a.len)

	get_ptr(a, idx)^ = get(a^, a.len - 1)
	a.len -= 1
}

clear :: proc "contextless" (a: $T/^Xarr($E, $SOA)) {
	a.len = 0
}

delete :: proc(a: $T/^Xarr($E, $SOA), allocator := context.allocator) {
	for i in 0 ..< builtin.len(a.chunks) {
		builtin.delete(get_chunk_slice(a^, i32(i)), allocator)
	}

	a^ = Xarr(E, SOA){}
}

Iterator :: struct($E: typeid, $SOA: bool) {
	xarr: ^Xarr(E, SOA),
	idx:  int,
}

iterator :: proc(a: $T/^Xarr($E, $SOA), start_idx := 0) -> Iterator(E, SOA) {
	return Iterator(E, SOA){xarr = a, idx = start_idx}
}

iterator_next :: proc(it: ^Iterator($E, $SOA)) -> (e: ^E, idx: int, ok: bool) {
	if it.idx >= it.xarr.len {
		return nil, it.idx, false
	}

	e = get_ptr(it.xarr, it.idx)
	idx = it.idx
	ok = true

	it.idx += 1
	return
}

Chunk_Iterator :: struct($E: typeid, $SOA: bool) {
	xarr:             ^Xarr(E, SOA),
	base_element_idx: int,
	chunk_idx:        i32,
}

chunk_iterator :: proc(a: $T/^Xarr($E, $SOA)) -> Chunk_Iterator(E, SOA) {
	return Chunk_Iterator(E, SOA){xarr = a}
}

chunk_iterator_next_scalar :: proc(
	it: ^Chunk_Iterator($E, false),
) -> (
	chunk: []E,
	base_element_idx: int,
	ok: bool,
) {
	if (it.xarr.allocated_chunks_mask & (u32(1) << u32(it.chunk_idx))) == 0 {
		return nil, 0, false
	}

	chunk = get_chunk_slice_scalar(it.xarr^, it.chunk_idx)
	// Limit the chunk to the length so user code doesn't have to worry about this
	base_element_idx = it.base_element_idx
	chunk = chunk[:min(builtin.len(chunk), it.xarr.len - base_element_idx)]
	ok = true

	base_element_idx += int(chunk_size(it.chunk_idx))
	it.chunk_idx += 1
	return
}

chunk_iterator_next_soa :: proc(
	it: ^Chunk_Iterator($E, true),
) -> (
	chunk: #soa[]E,
	base_element_idx: int,
	ok: bool,
) {
	if (it.xarr.allocated_chunks_mask & (u32(1) << u32(it.chunk_idx))) == 0 {
		return nil, 0, false
	}

	chunk = get_chunk_slice_soa(it.xarr^, it.chunk_idx)
	// Limit the chunk to the length so user code doesn't have to worry about this
	base_element_idx = it.base_element_idx
	chunk = chunk[:min(builtin.len(chunk), it.xarr.len - base_element_idx)]
	ok = true

	base_element_idx += int(chunk_size(it.chunk_idx))
	it.chunk_idx += 1
	return
}

chunk_iterator_next :: proc {
	chunk_iterator_next_scalar,
	chunk_iterator_next_soa,
}