mem :: import "@std/mem"

_probe_count func(context ?*mut anyopaque) void {
	if context |raw| {
		counts *mut usize :: ptr_cast(usize, raw)
		counts[0] += 1
	}
}

_probe_alloc func(context ?*mut anyopaque, _ usize, _ usize) ?*mut u8 {
	_probe_count(context)
	return none
}

_probe_realloc func(context ?*mut anyopaque, _ ?*mut u8, _ usize, _ usize, _ usize) ?*mut u8 {
	_probe_count(context)
	return none
}

_probe_free func(context ?*mut anyopaque, _ ?*mut u8, _ usize, _ usize) void {
	_probe_count(context)
}

_probe_vtable mem.AllocatorVTable :: mem.AllocatorVTable {
	alloc = _probe_alloc,
	realloc = _probe_realloc,
	free = _probe_free,
}

TaskList :: struct {
	ids ?[]mut i32
	priorities ?[]mut i32
	durations ?[]mut i32
	len usize
	capacity usize
	allocator mem.Allocator
}

task_list_init func(allocator mem.Allocator) TaskList {
	return TaskList {
		ids = none,
		priorities = none,
		durations = none,
		len = 0,
		capacity = 0,
		allocator = allocator,
	}
}

alloc_i32s func(allocator mem.Allocator, count usize) ?[]mut i32 {
	fallback [1]mut i32 = undefined
	failed bool = false
	values []mut i32 = mem.alloc(i32, allocator, count) catch |_| {
		failed = true
		yield (&fallback).ptr[..0]
	}
	if (failed) return none
	return values
}

free_i32s func(allocator mem.Allocator, values ?[]mut i32) void {
	if values |slice| {
		mem.raw_free(allocator, ptr_cast(u8, slice.ptr), slice.len * size_of(i32), align_of(i32))
	}
}

task_list_reserve func(list @mut TaskList, capacity usize) bool {
	if capacity <= list.capacity {
		return true
	}

	new_ids ?[]mut i32 = alloc_i32s(list.allocator, capacity)
	new_priorities ?[]mut i32 = alloc_i32s(list.allocator, capacity)
	new_durations ?[]mut i32 = alloc_i32s(list.allocator, capacity)

	if (new_ids and new_priorities and new_durations) |ids, priorities, durations| {
		if list.len > 0 {
			if (list.ids and list.priorities and list.durations) |old_ids, old_priorities, old_durations| {
				i usize = 0
				while i < list.len : i += 1 {
					ids[i] = old_ids[i]
					priorities[i] = old_priorities[i]
					durations[i] = old_durations[i]
				}
			} else {
				free_i32s(list.allocator, new_ids)
				free_i32s(list.allocator, new_priorities)
				free_i32s(list.allocator, new_durations)
				return false
			}
		}

		free_i32s(list.allocator, list.ids)
		free_i32s(list.allocator, list.priorities)
		free_i32s(list.allocator, list.durations)

		list.ids = new_ids
		list.priorities = new_priorities
		list.durations = new_durations
		list.capacity = capacity
		return true
	}

	free_i32s(list.allocator, new_ids)
	free_i32s(list.allocator, new_priorities)
	free_i32s(list.allocator, new_durations)
	return false
}

task_list_push func(list @mut TaskList, id i32, priority i32, duration i32) bool {
	if list.len == list.capacity {
		new_capacity usize = 2
		if list.capacity != 0 {
			new_capacity = list.capacity * 2
		}
		if task_list_reserve(list, new_capacity) == false {
			return false
		}
	}

	if (list.ids and list.priorities and list.durations) |ids, priorities, durations| {
		index usize = list.len
		ids[index] = id
		priorities[index] = priority
		durations[index] = duration
		list.len += 1
		return true
	}

	return false
}

task_score func(priority i32, duration i32) i32 {
	return priority * 10 - duration
}

task_list_best_id func(list @mut TaskList) i32 {
	if list.len == 0 {
		return -1
	}

	if (list.ids and list.priorities and list.durations) |ids, priorities, durations| {
		best_index usize = 0
		best_score i32 = task_score(priorities[0], durations[0])
		i usize = 1
		while i < list.len : i += 1 {
			score i32 = task_score(priorities[i], durations[i])
			if score > best_score {
				best_score = score
				best_index = i
			}
		}
		return ids[best_index]
	}

	return -1
}

task_list_total_duration func(list @mut TaskList) i32 {
	total i32 = 0
	if list.durations |durations| {
		i usize = 0
		while i < list.len : i += 1 {
			total += durations[i]
		}
	}
	return total
}

task_list_deinit func(list @mut TaskList) void {
	free_i32s(list.allocator, list.ids)
	free_i32s(list.allocator, list.priorities)
	free_i32s(list.allocator, list.durations)
	list.ids = none
	list.priorities = none
	list.durations = none
	list.len = 0
	list.capacity = 0
}

main func() i32 {
	if (size_of(mem.Allocator) != 16) return 31

	first_calls [1]mut usize = [0]
	second_calls [1]mut usize = [0]
	first_allocator mem.Allocator :: mem.Allocator {
		context = (&first_calls).ptr,
		vtable = &_probe_vtable,
	}
	second_allocator mem.Allocator :: mem.Allocator {
		context = (&second_calls).ptr,
		vtable = &_probe_vtable,
	}

	_ = mem.raw_alloc(first_allocator, 1, 1)
	_ = mem.raw_realloc(first_allocator, none, 0, 1, 1)
	mem.raw_free(first_allocator, none, 0, 1)
	_ = mem.raw_alloc(second_allocator, 1, 1)
	if (first_calls[0] != 3 or second_calls[0] != 1) return 32

	i32_fallback [1]mut i32 = undefined
	empty_failed bool = false
	empty []mut i32 = mem.alloc(i32, first_allocator, 0) catch |_| {
		empty_failed = true
		yield (&i32_fallback).ptr[..0]
	}
	if (empty_failed or empty.len != 0 or first_calls[0] != 3) return 34

	zero_sized_fallback [1]mut [0]u8 = undefined
	zero_sized_failed bool = false
	zero_sized []mut [0]u8 = mem.alloc([0]u8, first_allocator, 3) catch |_| {
		zero_sized_failed = true
		yield (&zero_sized_fallback).ptr[..0]
	}
	if (zero_sized_failed or zero_sized.len != 3 or first_calls[0] != 3) return 36
	_ = zero_sized[2]

	u64_fallback [1]mut u64 = undefined
	overflow_fallback_failed bool = false
	overflow_fallback []mut u64 = mem.alloc(u64, first_allocator, 0) catch |_| {
		overflow_fallback_failed = true
		yield (&u64_fallback).ptr[..0]
	}
	if (overflow_fallback_failed) return 37
	overflow_failed bool = false
	_ = mem.alloc(u64, first_allocator, max_value(usize)) catch |_| {
		overflow_failed = true
		yield overflow_fallback
	}
	if (overflow_failed == false or first_calls[0] != 3) return 40

	typed_failed bool = false
	typed []mut i32 = mem.alloc(i32, mem.c_allocator, 4) catch |_| {
		typed_failed = true
		yield (&i32_fallback).ptr[..0]
	}
	if (typed_failed) return 38
	defer mem.raw_free(mem.c_allocator, ptr_cast(u8, typed.ptr), typed.len * size_of(i32), align_of(i32))
	typed[0] = 10
	typed[3] = 20
	if (typed[0] + typed[3] != 30) return 39

	resized ?*mut u8 = mem.raw_realloc(mem.c_allocator, none, 0, 4, 1)
	if resized |bytes| {
		bytes[0] = 10
		bytes[1] = 20
		bytes[2] = 30
		bytes[3] = 40
	} else {
		return 20
	}

	grown ?*mut u8 = mem.raw_realloc(mem.c_allocator, resized, 4, 8, 1)
	if grown |bytes| {
		resized = grown
		if bytes[0] != 10 or bytes[1] != 20 or bytes[2] != 30 or bytes[3] != 40 {
			mem.raw_free(mem.c_allocator, grown, 8, 1)
			return 21
		}
	} else {
		mem.raw_free(mem.c_allocator, resized, 4, 1)
		return 22
	}

	shrunk ?*mut u8 = mem.raw_realloc(mem.c_allocator, resized, 8, 2, 1)
	if shrunk |bytes| {
		resized = shrunk
		if bytes[0] != 10 or bytes[1] != 20 {
			mem.raw_free(mem.c_allocator, shrunk, 2, 1)
			return 23
		}
	} else {
		mem.raw_free(mem.c_allocator, resized, 8, 1)
		return 24
	}

	invalid ?*mut u8 = mem.raw_realloc(mem.c_allocator, resized, 2, 4, 24)
	if invalid |memory| {
		mem.raw_free(mem.c_allocator, memory, 4, 24)
		mem.raw_free(mem.c_allocator, resized, 2, 1)
		return 25
	}
	if resized |bytes| {
		if bytes[0] != 10 or bytes[1] != 20 {
			mem.raw_free(mem.c_allocator, resized, 2, 1)
			return 26
		}
	}

	resized = mem.raw_realloc(mem.c_allocator, resized, 2, 0, 1)
	if resized |memory| {
		mem.raw_free(mem.c_allocator, memory, 0, 1)
		return 27
	}

	over_aligned ?*mut u8 = mem.raw_alloc(mem.c_allocator, 4, 32)
	if over_aligned |bytes| {
		bytes[0] = 11
		bytes[1] = 22
	} else {
		return 28
	}
	over_aligned_grown ?*mut u8 = mem.raw_realloc(mem.c_allocator, over_aligned, 4, 8, 32)
	if over_aligned_grown |bytes| {
		if bytes[0] != 11 or bytes[1] != 22 {
			mem.raw_free(mem.c_allocator, over_aligned_grown, 8, 32)
			return 29
		}
		mem.raw_free(mem.c_allocator, over_aligned_grown, 8, 32)
	} else {
		mem.raw_free(mem.c_allocator, over_aligned, 4, 32)
		return 30
	}

	zero_alignment ?*mut u8 = mem.raw_alloc(mem.c_allocator, 8, 0)
	if zero_alignment |memory| {
		mem.raw_free(mem.c_allocator, memory, 8, 0)
		return 1
	}

	bad_alignment ?*mut u8 = mem.raw_alloc(mem.c_allocator, 8, 24)
	if bad_alignment |memory| {
		mem.raw_free(mem.c_allocator, memory, 8, 24)
		return 2
	}

	aligned ?*mut u8 = mem.raw_alloc(mem.c_allocator, 64, 32)
	defer mem.raw_free(mem.c_allocator, aligned, 64, 32)
	if aligned |bytes| {
		bytes[0] = 1
		bytes[63] = 2
		if bytes[0] + bytes[63] != 3 {
			return 4
		}
	} else {
		return 3
	}

	tasks TaskList = task_list_init(mem.c_allocator)
	defer task_list_deinit(&tasks)

	if task_list_push(&tasks, 101, 3, 5) == false {
		return 5
	}
	if tasks.capacity != 2 {
		return 6
	}

	if task_list_push(&tasks, 202, 1, 4) == false {
		return 7
	}
	if tasks.capacity != 2 {
		return 8
	}

	if task_list_push(&tasks, 303, 4, 8) == false {
		return 9
	}
	if tasks.capacity != 4 {
		return 10
	}

	if tasks.len != 3 {
		return 11
	}

	if task_list_best_id(&tasks) != 303 {
		return 12
	}

	if task_list_total_duration(&tasks) != 17 {
		return 13
	}

	return 0
}