const std = @import("std");

const index = @import("index.zig");
const hashing = @import("hashing.zig");
const persistence = @import("persistence.zig");

const NUM_SHARDS = 64;
const TOTAL_BUCKETS = 1_048_576;
const BUCKETS_PER_SHARD = TOTAL_BUCKETS / NUM_SHARDS;

const Entry = struct {
    key: []const u8,
    value: []const u8,
    hash: u32,
    next: ?*Entry,
};

const Shard = struct {
    buckets: []?*Entry,
    rwlock: std.Thread.RwLock,
    arena: std.heap.ArenaAllocator,
    allocator: std.mem.Allocator,
};

var shards: [NUM_SHARDS]Shard = undefined;
var shards_initialized: bool = false;

var init_mutex: std.Thread.Mutex = .{};

fn getShardIndex(hash: u32) usize {
    return hash % NUM_SHARDS;
}

pub fn init() void {
    if (shards_initialized) return;

    init_mutex.lock();
    defer init_mutex.unlock();

    if (!shards_initialized) {
        for (&shards) |*shard| {
            shard.arena = std.heap.ArenaAllocator.init(std.heap.page_allocator);
            shard.allocator = shard.arena.allocator();
            shard.buckets = shard.allocator.alloc(?*Entry, BUCKETS_PER_SHARD) catch unreachable;
            @memset(shard.buckets, null);
            shard.rwlock = .{};
        }
        shards_initialized = true;
    }
}

pub fn restore(key: []const u8, value: []const u8) bool {
    const hash = hashing.hashKey(key);
    const shard_idx = getShardIndex(hash);

    var entry_key: []const u8 = undefined;
    {
        shards[shard_idx].rwlock.lock();
        defer shards[shard_idx].rwlock.unlock();

        const entry = writeVolatile(hash, key, value) orelse return false;
        entry_key = entry.key;
    }

    index.insert(entry_key);
    return true;
}

pub fn restoreDelete(key: []const u8) bool {
    const hash = hashing.hashKey(key);
    const shard_idx = getShardIndex(hash);

    const deleted = blk: {
        shards[shard_idx].rwlock.lock();
        defer shards[shard_idx].rwlock.unlock();
        break :blk deleteVolatile(hash, key);
    };

    if (deleted) {
        index.delete(key);
        return true;
    }
    return false;
}

pub fn writeVolatile(hash: u32, key: []const u8, value: []const u8) ?*Entry {
    const shard_idx = getShardIndex(hash);
    const bucketIdx = hash % shards[shard_idx].buckets.len;
    const alloc = shards[shard_idx].allocator;

    var current = shards[shard_idx].buckets[bucketIdx];
    while (current) |entry| {
        if (entry.hash == hash and std.mem.eql(u8, entry.key, key)) {
            alloc.free(entry.value);
            entry.value = alloc.dupe(u8, value) catch return null;
            return entry;
        }

        current = entry.next;
    }

    const newEntry = alloc.create(Entry) catch return null;
    errdefer alloc.destroy(newEntry);
    newEntry.* = Entry{
        .key = alloc.dupe(u8, key) catch return null,
        .value = alloc.dupe(u8, value) catch return null,
        .hash = hash, // Cache hash value
        .next = shards[shard_idx].buckets[bucketIdx],
    };

    shards[shard_idx].buckets[bucketIdx] = newEntry;
    return newEntry;
}

pub fn write(key: []const u8, value: []const u8) bool {
    const hash = hashing.hashKey(key);
    const shard_idx = getShardIndex(hash);

    var entry_key: []const u8 = undefined;
    var entry_value: []const u8 = undefined;
    {
        shards[shard_idx].rwlock.lock();
        defer shards[shard_idx].rwlock.unlock();

        const entry = writeVolatile(hash, key, value) orelse return false;
        entry_key = entry.key;
        entry_value = entry.value;
    }

    index.insert(entry_key);
    persistence.persist('W', entry_key, entry_value);
    return true;
}

pub fn read(key: []const u8) ?[]const u8 {
    if (!shards_initialized) return null;

    const hash = hashing.hashKey(key);
    const shard_idx = getShardIndex(hash);

    shards[shard_idx].rwlock.lockShared();
    defer shards[shard_idx].rwlock.unlockShared();

    var current = shards[shard_idx].buckets[hash % shards[shard_idx].buckets.len];
    while (current) |entry| {
        if (entry.hash == hash and std.mem.eql(u8, entry.key, key)) {
            return entry.value;
        }
        current = entry.next;
    }

    return null;
}

pub fn deleteVolatile(hash: u32, key: []const u8) bool {
    if (!shards_initialized) return false;
    const shard_idx = getShardIndex(hash);
    const bucketIdx = hash % shards[shard_idx].buckets.len;
    const alloc = shards[shard_idx].allocator;

    var current = shards[shard_idx].buckets[bucketIdx];
    var prev: ?*Entry = null;

    while (current) |entry| {
        if (entry.hash == hash and std.mem.eql(u8, entry.key, key)) {
            if (prev) |p| {
                p.next = entry.next;
            } else {
                shards[shard_idx].buckets[bucketIdx] = entry.next;
            }

            alloc.free(entry.key);
            alloc.free(entry.value);
            alloc.destroy(entry);
            return true;
        }

        prev = entry;
        current = entry.next;
    }

    return false;
}

pub fn delete(key: []const u8) bool {
    const hash = hashing.hashKey(key);
    const shard_idx = getShardIndex(hash);

    const deleted = blk: {
        shards[shard_idx].rwlock.lock();
        defer shards[shard_idx].rwlock.unlock();
        break :blk deleteVolatile(hash, key);
    };

    if (deleted) {
        index.delete(key);
        persistence.persist('D', key, "");
    }

    return deleted;
}