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

/// This array provides:
/// - O(1) insertion
/// - O(1) removal
/// - O(1) lookup
/// - Pointers and IDs stay stable when inserting and removing.
/// - Unique IDs for every inserted item (assuming that generation doesn't overflow)
///
/// For more details about this data structure checkout this podcast episode from Wookash:
/// * https://www.youtube.com/watch?v=ShSGHb65f3M
pub fn SlotMapType(Index: type, Generation: type, Value: type) type {
    assert(@typeInfo(Generation) == .int);
    assert(@typeInfo(Index) == .int);

    return struct {
        const Self = @This();

        // TODO: This struct will probably have paddings bytes inserted.
        // It's not ideal for a data structure like this which will be commonly used. 
        //
        // The `used` field could be packed into a bitset.
        // But that would introduce a second `std.ArrayList` which would make managing the memory a PITA.
        //
        // So in the end idk if it's worth it.
        pub const Slot = struct {
            value: Value,
            generation: Generation,
            next_hole: ?Index,
            used: bool
        };

        slots: std.ArrayList(Slot),
        first_hole: ?Index,
        last_hole: ?Index,
        hole_count: usize,

        const empty = Self{
            .slots = .empty,
            .first_hole = null,
            .last_hole = null,
            .hole_count = 0,
        };

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

            pub fn format(self: Id, writer: *std.Io.Writer) std.Io.Writer.Error!void {
                try writer.print("Id{{ {}, {} }}", .{ self.index, self.generation });
            }
        };

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

            pub fn next(self: *Iterator) ?Id {
                while (self.index < self.slot_map.slots.items.len) {
                    const index = self.index; 
                    const slot = self.slot_map.slots.items[index];
                    self.index += 1;

                    if (slot.used) {
                        return Id{
                            .index = @intCast(index),
                            .generation = slot.generation
                        };
                    }
                }

                return null;
            }
        };

        pub fn clearRetainingCapacity(self: *Self) void {
            const slots = self.slots;
            slots.clearRetainingCapacity();

            self.* = .init(slots);
        }

        fn insertHole(self: *Self, index: Index) void {
            if (self.last_hole) |last_hole| {
                self.slots.items[index].next_hole = last_hole;
                self.last_hole = index;
            } else {
                self.first_hole = index;
                self.last_hole = index;
            }
            self.hole_count += 1;
        }

        fn removeUnused(self: *Self) ?Index {
            if (self.first_hole) |first_hole| {
                self.first_hole = self.slots.items[first_hole].next_hole;
                if (self.first_hole == null) {
                    self.last_hole = null;
                }
                self.hole_count -= 1;
                return first_hole;
            }

            const capacity = @min(self.slots.capacity, std.math.maxInt(Index));
            if (self.slots.items.len < capacity) {
                const index: Index = @intCast(self.slots.items.len);
                self.slots.items.len += 1;
                self.slots.items[index] = Slot{
                    .value = undefined,
                    .generation = 0,
                    .next_hole = null,
                    .used = false
                };
                return index;
            }

            return null;
        }

        pub fn ensureUnusedCapacity(self: *Self, gpa: Allocator, additional_count: usize) Allocator.Error!void {
            if (additional_count > self.hole_count) {
                const new_capacity, const overflow = @addWithOverflow(
                    self.slots.capacity,
                    additional_count - self.hole_count
                );
                if (overflow != 0) return error.OutOfMemory;
                if (new_capacity >= std.math.maxInt(Index)) return error.OutOfMemory;

                try self.slots.ensureTotalCapacity(gpa, new_capacity);
            }
        }

        pub fn unusedCapacity(self: *Self) usize {
            const capacity = @min(self.slots.capacity, std.math.maxInt(Index));
            return capacity - self.slots.items.len + self.hole_count;
        }

        pub fn insertAssumeCapacity(self: *Self) Id {
            const index = self.removeUnused().?;
            const slot = &self.slots.items[index];
            assert(!slot.used);
            slot.used = true;

            return Id{
                .index = @intCast(index),
                .generation = slot.generation
            };
        }

        pub fn insert(self: *Self, gpa: Allocator) Allocator.Error!Id {
            try self.ensureUnusedCapacity(gpa, 1);
            return self.insertAssumeCapacity();
        }

        pub fn insertBounded(self: *Self) Allocator.Error!Id {
            if (self.unusedCapacity() == 0) {
                return error.OutOfMemory;
            }
            return self.insertAssumeCapacity();
        }

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

            const slot = self.slots.items[id.index];

            return slot.used and slot.generation == id.generation;
        }

        pub fn getAssumeExists(self: *Self, id: Id) *Value {
            assert(self.exists(id));
            return &self.slots.items[id.index].value;
        }

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

            const slot = &self.slots.items[id.index];

            slot.used = false;
            slot.generation +%= 1;
            self.insertHole(id.index);
        }

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

            self.removeAssumeExists(id);
            return true;
        }

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

        pub fn init(slots: std.ArrayList(Slot)) Self {
            var self: Self = .empty;
            self.slots = slots;
            return self;
        }

        pub fn deinit(self: *Self, gpa: Allocator) void {
            self.slots.deinit(gpa);
        }
    };
}

const TestMap = SlotMapType(u24, u8, void);

// TODO: Add more rigorous test for check if generation usage is nicely distributed.

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

    var map: TestMap = .empty;
    defer map.deinit(gpa);

    const id1 = try map.insert(gpa);
    try expect(map.exists(id1));
    try expect(map.remove(id1));
    try expect(!map.exists(id1));
    try expect(!map.remove(id1));

    const id2 = try map.insert(gpa);
    try expect(map.exists(id2));
    try expect(!map.exists(id1));
}

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

    var map: TestMap = .empty;
    defer map.deinit(gpa);

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

    // Artificially increase generation count
    map.slots.items[id1.index].generation = std.math.maxInt(@FieldType(TestMap.Id, "generation"));

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

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

    var map: TestMap = .empty;
    defer map.deinit(gpa);

    // Create array which has a hole
    const id1 = try map.insert(gpa);
    const id2 = try map.insert(gpa);
    const id3 = try map.insert(gpa);

    map.removeAssumeExists(id2);

    var iter = map.iterator();
    try expectEqual(id1, iter.next().?);
    try expectEqual(id3, iter.next().?);
    try expectEqual(null, iter.next());
}

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

    var map: TestMap = .empty;
    defer map.deinit(gpa);

    const id1 = try map.insert(gpa);
    try expect(map.exists(id1));
    map.clearRetainingCapacity();

    const id2 = try map.insert(gpa);
    try expect(map.exists(id2));

    try expectEqual(id1, id2);
}