const std = @import("std");
const tracy = @import("tracy");
const builtin = @import("builtin");
const assert = std.debug.assert;
const Allocator = std.mem.Allocator;

const Index = u24;
const Generation = u8;

pub fn GenerationalArrayList(Item: type) type {
    assert(@bitSizeOf(Generation) % 8 == 0);
    assert(@bitSizeOf(Index) % 8 == 0);

    return struct {
        const Self = @This();

        items: [*]Item,
        generations: [*]Generation,
        unused: [*]u8,

        len: u32,
        capacity: u32,

        count: u32,

        pub const empty = Self{
            .items = &[_]Item{},
            .generations = &[_]Generation{},
            .unused = &[_]u8{},
            .capacity = 0,
            .len = 0,
            .count = 0
        };

        pub const Id = packed struct {
            generation: Generation,
            index: Index,

            // TODO: Maybe `Id.Optional` type should be created to ensure .wrap() and .toOptional() 
            pub const none = Id{
                .generation = std.math.maxInt(Generation),
                .index = std.math.maxInt(Index),
            };

            pub fn format(self: Id, writer: *std.Io.Writer) std.Io.Writer.Error!void {
                if (self == Id.none) {
                    try writer.print("Id({s}){{ .none }}", .{ @typeName(Item) });
                } else {
                    try writer.print("Id({s}){{ {}, {} }}", .{ @typeName(Item), self.index, self.generation });
                }
            }

            pub fn asInt(self: Id) u32 {
                return @bitCast(self);
            }
        };

        pub const ItemWithId = struct {
            id: Id,
            item: *Item,
        };

        pub const Iterator = struct {
            array_list: *Self,
            index: Index,

            pub fn nextId(self: *Iterator) ?Id {
                while (self.index < self.array_list.len) {
                    const index = self.index;
                    self.index += 1;

                    // TODO: Inline the `byte_index` calculate for better speed.
                    // Probably not needed. Idk
                    if (self.array_list.isUnused(index)) {
                        continue;
                    }

                    return Id{
                        .index = @intCast(index),
                        .generation = self.array_list.generations[index]
                    };
                }

                return null;
            }

            pub fn nextItem(self: *Iterator) ?*Item {
                if (self.nextId()) |id| {
                    return &self.array_list.items[id.index];
                }

                return null;
            }

            pub fn next(self: *Iterator) ?ItemWithId {
                if (self.nextId()) |id| {
                    return ItemWithId{
                        .id = id,
                        .item = &self.array_list.items[id.index]
                    };
                }

                return null;
            }
        };

        pub const Metadata = extern struct {
            len: u32,
            count: u32
        };

        fn divCeilGeneration(num: u32) u32 {
            return std.math.divCeil(u32, num, @bitSizeOf(Generation)) catch unreachable;
        }

        fn divFloorGeneration(num: u32) u32 {
            return @divFloor(num, @bitSizeOf(Generation));
        }

        pub fn ensureTotalCapacityPrecise(self: *Self, allocator: Allocator, new_capacity: u32) !void {
            if (new_capacity > std.math.maxInt(Index)) {
                return error.OutOfIndexSpace;
            }

            // TODO: Shrinking is not supported
            assert(new_capacity >= self.capacity);

            const unused_bit_array_len = divCeilGeneration(self.capacity);
            const new_unused_bit_array_len = divCeilGeneration(new_capacity);

            // TODO: Handle allocation failure case
            const new_unused = try allocator.realloc(self.unused[0..unused_bit_array_len], new_unused_bit_array_len);
            const new_items = try allocator.realloc(self.items[0..self.capacity], new_capacity);
            const new_generations = try allocator.realloc(self.generations[0..self.capacity], new_capacity);

            self.unused = new_unused.ptr;
            self.items = new_items.ptr;
            self.generations = new_generations.ptr;
            self.capacity = new_capacity;
        }

        fn growCapacity(current: u32, minimum: u32) u32 {
            const init_capacity = @as(comptime_int, @max(1, std.atomic.cache_line / @sizeOf(Item)));

            var new = current;
            while (true) {
                new +|= new / 2 + init_capacity;
                if (new >= minimum) {
                    return new;
                }
            }
        }

        pub fn ensureTotalCapacity(self: *Self, allocator: Allocator, new_capacity: u32) !void {
            if (self.capacity >= new_capacity) return;

            const better_capacity = Self.growCapacity(self.capacity, new_capacity);
            try self.ensureTotalCapacityPrecise(allocator, better_capacity);
        }

        pub fn clearRetainingCapacity(self: *Self) void {
            self.count = 0;
            self.len = 0;
        }

        pub fn ensureUnusedCapacity(self: *Self, allocator: Allocator, unused_capacity: u32) !void {
            try self.ensureTotalCapacity(allocator, self.len + unused_capacity);
        }

        fn findFirstUnused(self: *Self) ?Index {
            for (0..divCeilGeneration(self.len)) |byte_index| {
                if (self.unused[byte_index] != 0) {
                    const found = @ctz(self.unused[byte_index]) + byte_index * @bitSizeOf(Generation);
                    if (found < self.len) {
                        return @intCast(found);
                    } else {
                        return null;
                    }
                }
            }

            return null;
        }

        fn markUnused(self: *Self, index: Index, unused: bool) void {
            assert(index < self.len);

            const byte_index = divFloorGeneration(index);
            const bit_index = @mod(index, @bitSizeOf(Generation));
            const bit_flag = @as(u8, 1) << @intCast(bit_index);
            if (unused) {
                self.unused[byte_index] |= bit_flag;
            } else {
                self.unused[byte_index] &= ~bit_flag;
            }
        }

        fn isUnused(self: *Self, index: Index) bool {
            assert(index < self.len);

            const byte_index = divFloorGeneration(index);
            const bit_index = @mod(index, @bitSizeOf(Generation));
            const bit_flag = @as(u8, 1) << @intCast(bit_index);
            return (self.unused[byte_index] & bit_flag) != 0;
        }

        pub fn insertUndefined(self: *Self, allocator: Allocator) !Id {
            var unused_index: Index = undefined;

            if (self.findFirstUnused()) |index| {
                unused_index = index;
            } else {
                try self.ensureUnusedCapacity(allocator, 1);

                unused_index = @intCast(self.len);
                self.len += 1;
                self.generations[unused_index] = 0;
            }

            self.markUnused(unused_index, false);
            self.count += 1;

            const id = Id{
                .index = @intCast(unused_index),
                .generation = self.generations[unused_index]
            };

            assert(id != Id.none);

            return id;
        }

        pub fn insert(self: *Self, allocator: Allocator, item: Item) !Id {
            const id = try self.insertUndefined(allocator);

            const new_item_ptr = self.getAssumeExists(id);
            new_item_ptr.* = item;

            return id;
        }

        pub fn exists(self: *Self, id: Id) bool {
            if (id.index >= self.len) {
                return false;
            }

            if (self.isUnused(id.index)) {
                return false;
            }

            if (self.generations[id.index] != id.generation) {
                return false;
            }

            return true;
        }

        pub fn removeAssumeExists(self: *Self, id: Id) void {
            assert(self.exists(id));

            self.markUnused(id.index, true);
            // TODO: Maybe a log should be shown when a wrap-around occurs?
            self.generations[id.index] +%= 1;
            self.count -= 1;
        }

        pub fn remove(self: *Self, id: Id) bool {
            if (!self.exists(id)) {
                return false;
            }

            self.removeAssumeExists(id);
            return true;
        }

        pub fn getAssumeExists(self: *Self, id: Id) *Item {
            assert(self.exists(id));

            return &self.items[id.index];
        }

        pub fn get(self: *Self, id: Id) ?*Item {
            if (self.exists(id)) {
                return self.getAssumeExists(id);
            } else {
                return null;
            }
        }

        pub fn iterator(self: *Self) Iterator {
            return Iterator{
                .array_list = self,
                .index = 0
            };
        }

        pub fn deinit(self: *Self, allocator: Allocator) void {
            allocator.free(self.unused[0..divCeilGeneration(self.capacity)]);
            allocator.free(self.generations[0..self.capacity]);
            allocator.free(self.items[0..self.capacity]);
        }

        pub fn getMetadata(self: *Self) Metadata {
            return Metadata{
                .len = self.len,
                .count = self.count
            };
        }

        pub fn write(self: *Self, writer: *std.Io.Writer, endian: std.builtin.Endian) !void {
            const zone = tracy.beginZone(@src(), .{ .name = "gen array list write" });
            defer zone.end();

            try writer.writeSliceEndian(Item, self.items[0..self.len], endian);
            try writer.writeSliceEndian(Generation, self.generations[0..self.len], endian);
            try writer.writeAll(self.unused[0..divCeilGeneration(self.len)]);
        }

        pub fn read(
            self: *Self,
            allocator: Allocator,
            reader: *std.Io.Reader,
            endian: std.builtin.Endian,
            metadata: Metadata
        ) !void {
            const zone = tracy.beginZone(@src(), .{ .name = "gen array list read" });
            defer zone.end();

            try self.ensureTotalCapacity(allocator, metadata.len);

            try reader.readSliceEndian(Item, self.items[0..metadata.len], endian);
            try reader.readSliceEndian(Generation, self.generations[0..metadata.len], endian);
            try reader.readSliceAll(self.unused[0..divCeilGeneration(metadata.len)]);

            self.len = metadata.len;
            self.count = metadata.count;
        }
    };
}

const TestArray = GenerationalArrayList(u32);

test "insert & remove" {
    const expect = std.testing.expect;
    const gpa = std.testing.allocator;

    var array_list: TestArray = .empty;
    defer array_list.deinit(gpa);

    const id1 = try array_list.insert(gpa, 10);
    try expect(array_list.exists(id1));
    try expect(array_list.remove(id1));
    try expect(!array_list.exists(id1));
    try expect(!array_list.remove(id1));

    const id2 = try array_list.insert(gpa, 10);
    try expect(array_list.exists(id2));
    try expect(!array_list.exists(id1));
    try expect(id1.index == id2.index);
}

test "generation wrap around" {
    const expectEqual = std.testing.expectEqual;
    const gpa = std.testing.allocator;

    var array_list: TestArray = .empty;
    defer array_list.deinit(gpa);

    // Grow array list so that at least 1 slot exists 
    const id1 = try array_list.insert(gpa, 10);
    array_list.removeAssumeExists(id1);

    // Artificially increase generation count
    array_list.generations[id1.index] = std.math.maxInt(Generation);

    // Check if generation wraps around
    const id2 = try array_list.insert(gpa, 10);
    array_list.removeAssumeExists(id2);
    try expectEqual(id1.index, id2.index);
    try expectEqual(0, array_list.generations[id1.index]);
}

test "iterator" {
    const expectEqual = std.testing.expectEqual;
    const gpa = std.testing.allocator;

    var array_list: TestArray = .empty;
    defer array_list.deinit(gpa);

    // Create array which has a hole
    const id1 = try array_list.insert(gpa, 1);
    const id2 = try array_list.insert(gpa, 2);
    const id3 = try array_list.insert(gpa, 3);

    array_list.removeAssumeExists(id2);

    var iter = array_list.iterator();
    try expectEqual(
        TestArray.ItemWithId{
            .id = id1,
            .item = array_list.getAssumeExists(id1)
        },
        iter.next().?
    );
    try expectEqual(
        TestArray.ItemWithId{
            .id = id3,
            .item = array_list.getAssumeExists(id3)
        },
        iter.next().?
    );
    try expectEqual(null, iter.next());
}

test "read & write" {
    const expectEqual = std.testing.expectEqual;
    const gpa = std.testing.allocator;

    var array_list1: TestArray = .empty;
    defer array_list1.deinit(gpa);

    var array_list2: TestArray = .empty;
    defer array_list2.deinit(gpa);

    const id1 = try array_list1.insert(gpa, 1);
    const id2 = try array_list1.insert(gpa, 2);
    const id3 = try array_list1.insert(gpa, 3);

    var buffer: [1024]u8 = undefined;
    var writer = std.Io.Writer.fixed(&buffer);
    const native_endian = builtin.cpu.arch.endian();

    try array_list1.write(&writer, native_endian);

    var reader = std.Io.Reader.fixed(writer.buffered());
    try array_list2.read(gpa, &reader, native_endian, array_list1.getMetadata());

    try expectEqual(array_list1.getAssumeExists(id1).*, array_list2.getAssumeExists(id1).*);
    try expectEqual(array_list1.getAssumeExists(id2).*, array_list2.getAssumeExists(id2).*);
    try expectEqual(array_list1.getAssumeExists(id3).*, array_list2.getAssumeExists(id3).*);
    try expectEqual(array_list1.count, array_list2.count);
}

test "clear retaining capacity" {
    const expect = std.testing.expect;
    const expectEqual = std.testing.expectEqual;
    const gpa = std.testing.allocator;

    var array_list: TestArray = .empty;
    defer array_list.deinit(gpa);

    const id1 = try array_list.insert(gpa, 10);
    try expect(array_list.exists(id1));
    array_list.clearRetainingCapacity();

    const id2 = try array_list.insert(gpa, 10);
    try expect(array_list.exists(id2));

    try expectEqual(id1, id2);
}