const std = @import("std");
const Handle = @import("handle.zig").Handle;

/// ----------------------------------------------------
/// ----------------------------------------------------
pub fn SlotMap(
    comptime K: type,
    comptime V: type,
) type {
    return struct {
        const Self = @This();

        /// ----------------------------------------------------
        /// ----------------------------------------------------
        pub const Entry = struct {
            key: K,
            value: V,
        };

        //
        // FIELDS
        //

        // --- VALUES ---
        values: std.ArrayList(Entry),
        lookup: std.AutoHashMap(K, Handle(V)),

        // --- LIFE TIMES ---
        generations: std.ArrayList(u32),
        free_list: std.ArrayList(u32),

        alloc: std.mem.Allocator,

        /// ----------------------------------------------------
        /// ----------------------------------------------------
        pub fn init(alloc: std.mem.Allocator) Self {
            return .{
                .values = .empty,
                .lookup = .init(alloc),

                .generations = .empty,
                .free_list = .empty,

                .alloc = alloc,
            };
        }

        /// ----------------------------------------------------
        /// ----------------------------------------------------
        pub fn deinit(self: *Self) void {
            self.values.deinit(self.alloc);
            self.lookup.deinit();
            self.generations.deinit(self.alloc);
            self.free_list.deinit(self.alloc);
        }

        /// ----------------------------------------------------
        /// ----------------------------------------------------
        pub fn put(
            self: *Self,
            key: K,
            value: V,
        ) error{ Duplicate, OutOfMemory }!Handle(V) {
            // --- PREVENT DUPS ---
            if (self.lookup.contains(key)) return error.Duplicate;

            // --- NEW GENERATION ---
            if (self.free_list.pop()) |idx| {
                self.values.items[@intCast(idx)] = .{
                    .key = key,
                    .value = value,
                };
                self.generations.items[@intCast(idx)] += 1;

                return .{
                    .idx = idx,
                    .gen = self.generations.items[idx],
                };
            }

            // --- METADATA ---
            const idx = self.values.items.len;
            const handle: Handle(V) = .{
                .idx = @intCast(idx),
                .gen = 1,
            };

            // --- NEW VALUE ---
            try self.values.append(self.alloc, .{
                .value = value,
                .key = key,
            });
            errdefer _ = self.values.swapRemove(idx);

            // --- NEW LOOKUP ---
            try self.lookup.put(key, handle);
            errdefer _ = self.lookup.remove(key);

            // --- NEW GENERATION ---
            try self.generations.append(self.alloc, 1);

            // --- RESULT ---
            return handle;
        }

        /// ----------------------------------------------------
        /// ----------------------------------------------------
        pub fn remove(
            self: *Self,
            handle: Handle(V),
        ) void {
            const idx: u32 = handle.idx;

            // --- CHECK IF STALE ---
            if (idx >= self.generations.items.len) return;
            if (self.generations.items[idx] != handle.gen) return;

            // --- REMOVE KEY ---
            const key = self.values.items[idx].key;
            _ = self.lookup.remove(key);

            // --- GENERATION ---
            self.generations.items[idx] +|= 1;
            self.free_list.append(self.alloc, idx) catch unreachable;
        }

        /// ----------------------------------------------------
        /// ----------------------------------------------------
        pub fn removeByKey(self: *Self, key: K) void {
            const idx = self.lookup.get(key) orelse return;
            _ = self.lookup.remove(key);
            self.generations.items[idx] +|= 1;
            self.free_list.append(self.alloc, idx) catch unreachable;
        }

        /// ----------------------------------------------------
        /// ----------------------------------------------------
        pub fn get(self: *Self, handle: Handle(V)) ?*V {
            if (handle.idx >= self.values.items.len) return null;
            if (self.generations.items[handle.idx] != handle.gen) return null;
            return &self.values.items[handle.idx].value;
        }

        /// ----------------------------------------------------
        /// ----------------------------------------------------
        pub fn getByKey(self: *Self, key: K) ?*V {
            const handle = self.lookup.get(key) orelse return null;
            return self.get(handle);
        }

        /// ----------------------------------------------------
        /// ----------------------------------------------------
        pub fn getHandle(self: *Self, key: K) ?Handle(V) {
            return self.lookup.get(key);
        }

        /// ----------------------------------------------------
        /// ----------------------------------------------------
        pub fn contains(self: *Self, key: K) bool {
            return self.lookup.contains(key);
        }

        /// ----------------------------------------------------
        /// ----------------------------------------------------
        pub fn count(self: *Self) u32 {
            return self.lookup.count();
        }
    };
}