const debug = @import("./debug.zig");

const io = @import("./io.zig");

const math = @import("./math.zig");

const table = @import("./table.zig");

pub const Bucketed = struct {
	base_allocator: io.Allocator,
	slab_table: SlabTable,

	const Slab = struct {
		count: usize = 0,
		buffer: []u8 = &.{},
		erased: []usize = &.{},

		fn create(self: *Slab, allocator: io.Allocator) ?[]u8 {
			if (self.count == self.erased.len) {
				const buffer = io.allocate_many(allocator, u8, math.max(1, self.buffer.len * 2)) orelse return null;

				errdefer io.deallocate(allocator, buffer);

				const erased = io.allocate_many(allocator, usize, math.max(1, self.erased.len * 2)) orelse return null;

				errdefer io.deallocate(allocator, erased);

				self.buffer = buffer;
				self.erased = erased;
			}

			return null;
		}

		fn destroy(self: *Slab) void {
			_ = self;
		}
	};

	const SlabTable = table.Hashed(table.unsigned_key(@bitSizeOf(usize)), *Slab);

	fn acquire_slab(self: *Bucketed, slab_element_size: usize) ?*Slab {
		if (slab_element_size == 0) return null;

		return self.slab_table.lookup(slab_element_size) orelse create_slab: {
			const allocated_slab = io.allocate_one(self.base_allocator, Slab);

			errdefer io.deallocate(self.base_allocator, allocated_slab);

			allocated_slab.* = .{.size = slab_element_size};

			debug.assert(self.size_buckets.insert(slab_element_size, allocated_slab) catch return null);

			break: create_slab allocated_slab;
		};
	}

	pub fn as_allocator(self: *Bucketed) io.Allocator {
		return io.Allocator.bind(self, reallocate);
	}

	pub fn deinit(self: *Bucketed) void {
		var slab_iterator = SlabTable.Iterator{.table = self.slab_table};

		while (slab_iterator.next()) |slab| {
			slab.free(self.base_allocator);
		}

		self.size_buckets.free(self.base_allocator);
	}

	pub fn init(base_allocator: io.Allocator) io.AllocationError!Bucketed {
		return Bucketed{
			.base_allocator = base_allocator,
			.size_buckets = &.{},
		};
	}

	pub fn owns(self: Bucketed, memory: []const u8) bool {
		return io.overlaps(memory.ptr, (self.slab_table.lookup(memory.len) orelse return false).buffer);
	}

	pub fn reallocate(self: *Bucketed, options: io.AllocationOptions) ?[]u8 {
		const origin_slab = self.acquire_slab(options.size) orelse return null;
		const existing_allocation = options.allocation orelse return origin_slab.create(self.base_allocator);

		defer origin_slab.destroy(existing_allocation);

		const target_slab = self.acquire_slab(existing_allocation.len) orelse return null;
		const updated_allocation = target_slab.create(existing_allocation.len);

		io.copy(updated_allocation, existing_allocation);

		return updated_allocation;
	}
};