package spanpool

import "base:builtin"
import "core:fmt"
import "core:slice"

Span :: struct {
	first, len: i32,
}

Generation_Index :: i32

Span_Pool :: struct($E: typeid) {
	elems:       [dynamic]E,
	generations: [dynamic]Generation_Index,
	free_spans:  [dynamic]Span,
}

Handle :: struct {
	first: i32,
	len:   i32,
	gen:   i32,
}

destroy_spanpool :: proc(s: ^$T/Span_Pool($E)) {
	delete(s.elems)
	delete(s.generations)
	delete(s.free_spans)
}

// TODO: use everywhere
is_handle_valid :: proc(s: ^$T/Span_Pool($E), handle: Handle) -> bool {
	return(
		handle.gen != 0 &&
		int(handle.first) < len(s.elems) &&
		int(handle.first + handle.len) <= len(s.elems) &&
		s.generations[handle.first] == handle.gen \
	)
}

resolve_slice :: proc(s: ^$T/Span_Pool($E), handle: Handle, loc := #caller_location) -> []E {
	if !is_handle_valid(s, handle) {
		return {}
	}

	return s.elems[handle.first:handle.first + handle.len]
}

resolve_single :: proc(s: ^$T/Span_Pool($E), handle: Handle) -> ^E {
	assert(int(handle.first + handle.len) <= len(s.elems))
	assert(handle.len == 1)
	assert(s.generations[handle.first] == handle.gen)

	return &s.elems[handle.first]
}

allocate_elems :: proc(s: ^$T/Span_Pool($E), elems: ..E) -> (handle: Handle) {
	handle = allocate_num(s, i32(len(elems)))

	for i in 0 ..< handle.len {
		s.elems[i + handle.first] = elems[i]
	}

	return
}

allocate_num :: proc(s: ^$T/Span_Pool($E), count: i32) -> (handle: Handle) {
	handle.len = count

	maybe_existing_span: Maybe(i32)
	for span, i in s.free_spans {
		if span.len >= count {
			maybe_existing_span = i32(i)
			break
		}
	}

	existing_span, ok := maybe_existing_span.?
	if ok {
		span := s.free_spans[existing_span]
		handle.first = span.first
		new_len := span.len - count

		if new_len == 0 {
			ordered_remove(&s.free_spans, existing_span)
		} else {
			s.free_spans[existing_span].first += count
			s.free_spans[existing_span].len = new_len
		}
	} else {
		handle.first = i32(len(s.elems))
		new_len := len(s.elems) + int(count)
		resize(&s.elems, new_len)
		resize(&s.generations, new_len)
	}

	// Now figure out the generation index
	max_gen := i32(0)
	for i in handle.first ..< handle.first + handle.len {
		max_gen = max(max_gen, s.generations[i])
	}

	handle.gen = max_gen + 1
	for i in handle.first ..< handle.first + handle.len {
		s.generations[i] = handle.gen
	}

	return
}

free :: proc(s: ^$T/Span_Pool($E), handle: Handle, loc := #caller_location) {
	fmt.assertf(
		int(handle.first + handle.len) <= len(s.elems),
		"invalid handle %v",
		handle,
		loc = loc,
	)

	for i in 0 ..< handle.len {
		fmt.assertf(
			s.generations[handle.first + i] == handle.gen,
			"wrong generation number %v at index %v for handle %v",
			s.generations[handle.first + i],
			i,
			handle,
			loc = loc,
		)
		s.generations[handle.first + i] += 1
	}
	append(&s.free_spans, Span{first = handle.first, len = handle.len})
}

reconcile :: proc(s: ^$T/Span_Pool($E)) {
	slice.sort_by_key(s.free_spans[:], proc(span: Span) -> i32 {
		return span.first
	})

	for i := len(s.free_spans) - 2; i >= 0; i -= 1 {
		cur_span, next_span := &s.free_spans[i], &s.free_spans[i + 1]

		if (cur_span.first + cur_span.len == next_span.first) {
			// Merge
			cur_span.len += next_span.len
			unordered_remove(&s.free_spans, i + 1)
		}
	}
}

copy :: proc(dst: ^Span_Pool($E), src: Span_Pool(E)) {
	resize(&dst.elems, len(src.elems))
	resize(&dst.generations, len(src.generations))
	resize(&dst.free_spans, len(src.free_spans))

	builtin.copy(dst.elems[:], src.elems[:])
	builtin.copy(dst.generations[:], src.generations[:])
	builtin.copy(dst.free_spans[:], src.free_spans[:])
}