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

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

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

        //
        // FIELDS
        //

        // --- VALUES ---
        slot_map: SlotMap(Entry),
        lookup: std.AutoHashMap(K, Handle(V)),

        // --- MEMORY ---
        alloc: std.mem.Allocator,

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

                .alloc = alloc,
            };
        }

        /// ----------------------------------------------------
        /// ----------------------------------------------------
        pub fn deinit(self: *Self) void {
            self.slot_map.deinit();
            self.lookup.deinit();
        }

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

            // --- RESULT ---
            const handle = try self.slot_map.put(.{ .value = value, .key = key });
            errdefer self.slot_map.remove(handle);

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

            return handle;
        }

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

            // --- REMOVE FROM SLOTMAP ---
            self.slot_map.remove(handle);

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

        /// ----------------------------------------------------
        /// ----------------------------------------------------
        pub fn removeByKey(self: *Self, key: K) void {
            const handle = self.lookup.get(key) orelse return;
            _ = self.lookup.remove(key);
            self.slot_map.remove(handle);
        }

        /// ----------------------------------------------------
        /// ----------------------------------------------------
        pub fn get(self: *Self, handle: Handle(V)) ?*V {
            const entry = if (self.slot_map.get(handle)) |v| v else return null;
            return &entry.value;
        }

        /// ----------------------------------------------------
        /// ----------------------------------------------------
        pub fn getOrPut(
            self: *Self,
            key: K,
            value: V,
        ) !*V {
            if (self.getByKey(key)) |v| return v;
            const handle = try self.put(key, value);
            return self.get(handle).?;
        }

        /// ----------------------------------------------------
        /// ----------------------------------------------------
        pub fn getOrPutHandle(
            self: *Self,
            key: K,
            value: V,
        ) !*Handle(V) {
            if (self.getHandle(key)) |h| return h;
            return try self.put(key, 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();
        }

        /// ----------------------------------------------------
        /// ----------------------------------------------------
        pub fn items(self: *Self) []Entry {
            return self.slot_map.values.items;
        }
    };
}