const coral = @import("./coral.zig");

const std = @import("std");

pub fn Hashed(comptime Key: type, comptime Value: type, comptime traits: Traits(Key)) type {
    const load_max = 0.75;
    const max_int = std.math.maxInt(usize);

    return struct {
        entry_map: []?Entry,
        len: usize,

        pub const Entry = struct {
            key: Key,
            value: Value,

            fn writeInto(self: Entry, table: *Self) ?*Entry {
                const hash_max = @min(max_int, table.entry_map.len);
                var hashed_key = traits.hash(self.key) % hash_max;
                var iterations = @as(usize, 0);

                while (true) : (iterations += 1) {
                    std.debug.assert(iterations < table.entry_map.len);

                    const table_entry = &(table.entry_map[hashed_key] orelse {
                        table.entry_map[hashed_key] = .{
                            .key = self.key,
                            .value = self.value,
                        };

                        table.len += 1;

                        return &table.entry_map[hashed_key].?;
                    });

                    if (traits.are_equal(table_entry.key, self.key)) {
                        return null;
                    }

                    hashed_key = (hashed_key +% 1) % hash_max;
                }
            }
        };

        pub const Entries = struct {
            table: *const Self,
            iterations: usize,

            pub fn next(self: *Entries) ?*Entry {
                while (self.iterations < self.table.entry_map.len) {
                    defer self.iterations += 1;

                    if (self.table.entry_map[self.iterations]) |*entry| {
                        return entry;
                    }
                }

                return null;
            }
        };

        const Self = @This();

        pub const empty = Self{
            .entry_map = &.{},
            .len = 0,
        };

        pub fn entries(self: *const Self) Entries {
            return .{
                .table = self,
                .iterations = 0,
            };
        }

        pub fn isEmpty(self: Self) bool {
            return self.len == 0;
        }

        pub fn remove(self: *Self, key: Key) ?Entry {
            const hash_max = @min(max_int, self.entry_map.len);
            var hashed_key = key.hash() % hash_max;

            while (true) {
                const entry = &(self.entry_map[hashed_key] orelse continue);

                if (self.keys_equal(entry.key, key)) {
                    const original_entry = entry.*;

                    self.entry_map[hashed_key] = null;

                    return original_entry;
                }

                hashed_key = (hashed_key +% 1) % hash_max;
            }
        }

        pub fn replace(self: *Self, key: Key, value: Value) std.mem.Allocator.Error!?Entry {
            try self.rehash(load_max);

            std.debug.assert(self.entry_map.len > self.len);

            {
                const hash_max = @min(max_int, self.entry_map.len);
                var hashed_key = traits.hash(key) % hash_max;

                while (true) {
                    const entry = &(self.entry_map[hashed_key] orelse {
                        self.entry_map[hashed_key] = .{
                            .key = key,
                            .value = value,
                        };

                        self.len += 1;

                        return null;
                    });

                    if (traits.are_equal(key, entry.key)) {
                        const original_entry = entry.*;

                        entry.* = .{
                            .key = key,
                            .value = value,
                        };

                        return original_entry;
                    }

                    hashed_key = (hashed_key +% 1) % hash_max;
                }
            }
        }

        pub fn clear(self: *Self) void {
            for (self.entry_map) |*entry| {
                entry.* = null;
            }

            self.len = 0;
        }

        pub fn deinit(self: *Self) void {
            if (self.entry_map.len == 0) {
                return;
            }

            coral.heap.allocator.free(self.entry_map);

            self.* = undefined;
        }

        pub fn get(self: Self, key: Key) ?*Value {
            if (self.len == 0) {
                return null;
            }

            const hash_max = @min(max_int, self.entry_map.len);
            var hashed_key = traits.hash(key) % hash_max;
            var iterations = @as(usize, 0);

            while (iterations < self.len) : (iterations += 1) {
                const entry = &(self.entry_map[hashed_key] orelse return null);

                if (traits.are_equal(entry.key, key)) {
                    return &entry.value;
                }

                hashed_key = (hashed_key +% 1) % hash_max;
            }

            return null;
        }

        pub fn emplace(self: *Self, key: Key, value: Value) std.mem.Allocator.Error!switch (Value) {
            void => bool,
            else => ?*Value,
        } {
            try self.rehash(load_max);

            std.debug.assert(self.entry_map.len > self.len);

            const entry = Entry{
                .key = key,
                .value = value,
            };

            if (entry.writeInto(self)) |written_entry| {
                return switch (Value) {
                    void => true,
                    else => &written_entry.value,
                };
            }

            return switch (Value) {
                void => false,
                else => null,
            };
        }

        pub fn loadFactor(self: Self) f32 {
            return if (self.entry_map.len == 0) 1 else @as(f32, @floatFromInt(self.len)) / @as(f32, @floatFromInt(self.entry_map.len));
        }

        pub fn rehash(self: *Self, max_load: f32) std.mem.Allocator.Error!void {
            if (self.loadFactor() <= max_load) {
                return;
            }

            var table = empty;

            errdefer table.deinit();

            table.entry_map = allocate: {
                const min_len = @max(1, self.len);
                const table_size = min_len * 2;
                const zeroed_entry_map = try coral.heap.allocator.alloc(?Entry, table_size);

                errdefer {
                    coral.heap.allocator.free(zeroed_entry_map);
                }

                for (zeroed_entry_map) |*entry| {
                    entry.* = null;
                }

                break :allocate zeroed_entry_map;
            };

            for (self.entry_map) |maybe_entry| {
                if (maybe_entry) |entry| {
                    std.debug.assert(entry.writeInto(&table) != null);
                }
            }

            self.deinit();

            self.* = table;
        }
    };
}

pub fn Traits(comptime Key: type) type {
    return struct {
        are_equal: fn (Key, Key) bool,
        hash: fn (Key) usize,
    };
}

pub fn enumTraits(comptime Enum: type) Traits(Enum) {
    const enums = struct {
        fn are_equal(a: Enum, b: Enum) bool {
            return a == b;
        }

        fn hash(value: Enum) usize {
            return @intFromEnum(value) % std.math.maxInt(usize);
        }
    };

    return .{
        .are_equal = enums.are_equal,
        .hash = enums.hash,
    };
}

pub fn ptrTraits(comptime Ptr: type) Traits(Ptr) {
    const pointers = struct {
        fn are_equal(a: Ptr, b: Ptr) bool {
            return a == b;
        }

        fn hash(value: Ptr) usize {
            return @intFromPtr(value);
        }
    };

    return switch (@typeInfo(Ptr)) {
        .pointer => .{
            .are_equal = pointers.are_equal,
            .hash = pointers.hash,
        },

        else => @compileError(std.fmt.comptimePrint("parameter `Ptr` must be a pointer type, not {s}", .{
            @typeName(Ptr),
        })),
    };
}

pub const string_traits = init: {
    const strings = struct {
        fn are_equal(a: []const u8, b: []const u8) bool {
            return std.mem.eql(u8, a, b);
        }

        fn hash(value: []const u8) usize {
            return coral.hashes.djb2(usize, value);
        }
    };

    break :init Traits([]const u8){
        .are_equal = strings.are_equal,
        .hash = strings.hash,
    };
};

pub const usize_traits = init: {
    const usizes = struct {
        fn are_equal(a: usize, b: usize) bool {
            return a == b;
        }

        fn hash(value: usize) usize {
            return value;
        }
    };

    break :init Traits(usize){
        .are_equal = usizes.are_equal,
        .hash = usizes.hash,
    };
};