const std = @import("std");
const mem = std.mem;
const math = std.math;

const warn = std.debug.warn;
const assert = std.debug.assert;

const ArrayList = std.ArrayList;

fn Bucket(comptime T: type, comptime items_per_bucket: i32) type {
    return struct {
        const Self = @This();

        data: [items_per_bucket]T,
        occupied: [items_per_bucket]bool,
        index: i32,
        count: i32,
    };
}

const BucketLocator = struct {
    bucket_index: i32,
    slot_index: usize,
};

fn indexOf(comptime T: type, haystack: var, needle: T) ?usize {
    var itr = haystack.iterator();
    var i: usize = 0;
    while (itr.next()) |e| : (i += 1) {
        if (e == needle)  return i;
    }
    return null;
}

// TODO: Add shrink to free empty buckets.
fn BucketArray(comptime T: type, comptime items_per_bucket: i32) type {
    return struct {
        const MyBucket = Bucket(T, items_per_bucket);
        const Self = @This();

        allocator: *mem.Allocator,
        count: i64,

        all_buckets: ArrayList(*MyBucket),
        unfull_buckets: ArrayList(*MyBucket),


        pub fn init(allocator: *mem.Allocator) Self {
            var a = Self {
                .all_buckets = ArrayList(*MyBucket).init(allocator),
                .unfull_buckets = ArrayList(*MyBucket).init(allocator),
                .count = 0,
                .allocator = allocator,
            };
            return a;
        }

        pub fn deinit(self: *Self) void {
            var itr = self.all_buckets.iterator();
            while (itr.next()) |bucket| {
                self.allocator.destroy(bucket);
            }
            self.unfull_buckets.shrink(0);
            self.all_buckets.shrink(0);
        }

        fn addBucket(self: *Self) error{OutOfMemory}!void {
            var b = try self.allocator.create(MyBucket);
            b.index = @intCast(i32, self.all_buckets.len);
            b.count = 0;
            assert(b.index == @intCast(i32, self.all_buckets.len)); // NOTE: Asserts if overflowed.

            try self.all_buckets.append(b);
            try self.unfull_buckets.append(b);
        }

        pub fn append(self: *Self, item: T) error{OutOfMemory}!BucketLocator {
            if (self.unfull_buckets.len == 0)  try addBucket(self);
            assert(self.unfull_buckets.len != 0);

            var bucket = self.unfull_buckets.at(0);
            var index: ?usize = null;
            for (bucket.occupied) |occupied, i| {
                if (occupied) { // TODO: FIXME: Why do we need to have this branch in order for this code to work at all?
                    continue;
                } else {
                    index = i;
                    break;
                }
                // if (!occupied) {
                //     index = i;
                //     break;
                // }
            }
            assert(index != null); // TODO: Asserts.

            bucket.occupied[index.?] = true;
            bucket.count += 1;
            assert(bucket.count <= items_per_bucket);

            self.count += 1;

            if (bucket.count == items_per_bucket) {
                var itr = self.unfull_buckets.iterator();
                var i: usize = 0;
                var removed = false;
                while (itr.next()) |value| : (i += 1) {
                    if (value == bucket) {
                        _ = self.unfull_buckets.swapRemove(i);
                        removed = true;
                        break;
                    }
                }
                assert(removed);
            }

            var loc = BucketLocator{
                .bucket_index = bucket.index,
                .slot_index = index.?,
            };

            bucket.data[index.?] = item;
            return loc;
        }

        pub fn remove(self: *Self, loc: BucketLocator) !T {
            var bucket = self.all_buckets.at(@intCast(usize, loc.bucket_index));
            assert(bucket.occupied[loc.slot_index] == true);

            const wasFull = (bucket.count == items_per_bucket);

            bucket.occupied[loc.slot_index] = false;
            bucket.count -= 1;
            self.count -= 1;

            if (wasFull) {
                assert(indexOf(*MyBucket, self.unfull_buckets, bucket) == null);
                try self.unfull_buckets.append(bucket);
            }

            return bucket.data[loc.slot_index];
        }

        pub fn at(self: *Self, loc: BucketLocator) T {
            var bucket = self.all_buckets.at(@intCast(usize, loc.bucket_index));
            assert(bucket.occupied[loc.slot_index] == true);
            var result = bucket.data[@intCast(usize, loc.slot_index)];
            return result;
        }

        pub fn iterator(self: *Self) BucketIterator {
            return BucketIterator{
                .array = self,
                .locator = BucketLocator{.bucket_index = 0, .slot_index = 0},
                .index = -1,
            };
        }

        const BucketIterator = struct {
            array: *Self,
            index: i64,
            locator: BucketLocator,

            fn next(self: *BucketIterator) ?T {
                if (self.index >= self.array.count)  return null;

                self.index += 1;
                var loc = self.locator;

                if (self.index > 0) {
                    while (true) {
                        loc.slot_index += 1;
                        if (loc.slot_index >= @intCast(usize, items_per_bucket)) {
                            loc.slot_index = 0;
                            loc.bucket_index += 1;
                        }

                        if (@intCast(usize, loc.bucket_index) >= self.array.all_buckets.len)  return null;

                        var bucket = self.array.all_buckets.at(@intCast(usize, loc.bucket_index));
                        if (bucket.occupied[@intCast(usize, loc.slot_index)])  break;
                    }
                    self.locator = loc;
                }

                return self.array.all_buckets.at(@intCast(usize, loc.bucket_index)).data[@intCast(usize, loc.slot_index)];
            }
        };
    };
}

pub fn main() !void {
    var arena = std.heap.ArenaAllocator.init(&std.heap.DirectAllocator.init().allocator);
    defer arena.deinit();

    var ba = BucketArray(i32, 4).init(&arena.allocator);
    defer ba.deinit();
    assert(ba.count == 0);

    {
        var i: i32 = 1;
        while (i <= 1024) : (i *= 2) {
            _ = try ba.append(i);
            warn("add {}\n",  i);
        }
        assert(ba.count == math.log2(i));
    }
    const countAfterAdding = ba.count;

    {
        var itr = ba.iterator();
        var j: i32 = 0;
        warn("{{ ");
        while (itr.next()) |value| : (j += 1) {
            if (j > 0) warn(", ");
            warn("{}", value);
            assert(value == math.pow(i32, 2, j));
        }
        warn(" }}\n");
    }

    {
        const removed = try ba.remove(BucketLocator{.bucket_index = 0, .slot_index = 2});
        assert(removed == 4);
        assert(ba.count == countAfterAdding-1);
    }

    {
        var itr = ba.iterator();
        var i: i32 = 0;
        warn("{{ ");
        while (itr.next()) |value| : (i += 1) {
            if (i > 0) warn(", ");
            warn("{}", value);

            assert(value != 4); // Since we removed it.
        }
        warn(" }}\n");
    }
}