# hashtable - Open-addressing hash table library
# Linear probing for collision resolution, supports auto-resizing and tombstone deletion
# Key type fixed as str (i8*), value type is void* (generic pointer)
import t, c
import string
import memhub
import stdio
from stdint import *

# Empty slot marker (hash == 0)
define _HT_EMPTY = 0
# Tombstone slot marker (hash == 1)
define _HT_TOMBSTONE = 1
# Default initial capacity (must be power of two)
define _HT_INIT_CAP = 16
# Slot memory layout: hash(8) + key pointer(8) + value pointer(8) = 24 bytes
define _HT_SLOT_SIZE = 24
# FNV-1a 64-bit constants
_FNV_OFFSET: uint64 = 0xcbf29ce484222325
_FNV_PRIME: uint64 = 0x100000001b3


def _fnv1a_hash(key: str) -> uint64:
    """FNV-1a 64-bit hash function."""
    h: uint64 = _FNV_OFFSET
    k: str = key
    while k[0] != 0:
        h = h ^ uint64(k[0])
        h = h * _FNV_PRIME
        k += 1
    # Ensure h >= 2 (avoids empty slot marker 0 and tombstone marker 1)
    if h < 2:
        h = 2
    return h


class HashTable:
    """Open-addressing hash table using linear probing for collision resolution.

    Key type fixed as str (i8*), value type is void* (generic pointer).
    Memory is managed via the mbuddy allocator. Automatically resizes when load factor exceeds 0.75.

    Usage:
        mb: memhub.MemBuddy = memhub.MemBuddy(arena, arena_size)
        ht: HashTable = HashTable(mb)
        ht.set_int("x", 42)
        p: t.CPtr = ht["x"]
        v: int = c.Deref((int | t.CPtr)(p))    # int value stored as i32; read it back with an int* view
    """
    __slots__:      t.CPtr              # Base address of slot array
    __count__:      sizet = 0                  # Number of occupied entries
    __capacity__:   sizet = 16                 # Total slot count (power of two)
    __tombstones__: sizet = 0                  # Tombstone entry count
    __mbuddy__:     memhub.MemBuddy*
    __iter_index__: sizet = 0

    def __new__(self, mb: memhub.MemBuddy*) -> t.CPtr:
        # Heap-allocate HashTable object (6 fields × 8 bytes = 48 bytes)
        return mb.alloc(48)

    def __init__(self, mb: memhub.MemBuddy*):
        self.__mbuddy__ = mb
        self.__capacity__ = _HT_INIT_CAP
        self.__count__ = 0
        self.__tombstones__ = 0
        self.__slots__ = mb.alloc(self.__capacity__ * _HT_SLOT_SIZE)
        # Initialize all slots to empty (hash = 0)
        string.memset(self.__slots__, 0, self.__capacity__ * _HT_SLOT_SIZE)

    def _slot_hash(self, idx: sizet) -> uint64:
        """Read hash value from target slot."""
        ptr: uint64* = t.CPtr(uint64(self.__slots__) + idx * _HT_SLOT_SIZE)
        return ptr[0]

    def _slot_key(self, idx: sizet) -> str:
        """Read key pointer from target slot."""
        ptr: str* = (str*)(uint64(self.__slots__) + idx * _HT_SLOT_SIZE + 8)
        return ptr[0]

    def _slot_value(self, idx: sizet) -> t.CPtr:
        """Read value pointer from target slot."""
        # The value slot holds a pointer, so the annotation must be "pointer to pointer": X|t.CPtr merges as Ptr(resolve(X)), so t.CUInt64T|t.CPtr|t.CPtr = Ptr(Ptr(i64)) = i64**, and subscripting yields i64* (a pointer).
        # Writing only t.CUInt64T|t.CPtr (= i64*, for hash integers) would make subscripting yield a scalar i64, hitting "implicit int-to-pointer conversion forbidden".
        ptr: uint64** = (uint64**)(uint64(self.__slots__) + idx * _HT_SLOT_SIZE + 16)
        return ptr[0]

    def _set_slot(self, idx: sizet, h: uint64, key: str, val: t.CPtr):
        """Assign three fields of the target slot."""
        h_ptr: uint64* = t.CPtr(uint64(self.__slots__) + idx * _HT_SLOT_SIZE)
        h_ptr[0] = h
        k_ptr: str* = (str*)(uint64(self.__slots__) + idx * _HT_SLOT_SIZE + 8)
        k_ptr[0] = key
        # Same as _slot_value: the value slot is annotated as "pointer to pointer" (i64**) so the element is a pointer
        v_ptr: uint64** = (uint64**)(uint64(self.__slots__) + idx * _HT_SLOT_SIZE + 16)
        v_ptr[0] = val

    def _find_slot(self, key: str, h: uint64) -> sizet:
        """Locate slot index for key insertion or lookup.

        Returns: Index of matching key if found; otherwise returns index of first usable empty/tombstone slot.
        """
        mask: sizet = self.__capacity__ - 1
        idx: sizet = sizet(h) & mask
        first_tombstone: sizet = self.__capacity__
        for i in range(self.__capacity__):
            sh: uint64 = self._slot_hash(idx)
            if sh == _HT_EMPTY:
                if first_tombstone < self.__capacity__:
                    return first_tombstone
                return idx
            if sh == _HT_TOMBSTONE:
                if first_tombstone >= self.__capacity__:
                    first_tombstone = idx
            else:
                if sh == h:
                    sk: str = self._slot_key(idx)
                    if string.strcmp(sk, key) == 0:
                        return idx
            idx = (idx + 1) & mask
        if first_tombstone < self.__capacity__:
            return first_tombstone
        return self.__capacity__

    def _resize(self, new_cap: sizet):
        """Resize table and rehash all entries."""
        old_slots: t.CPtr = self.__slots__
        old_cap: sizet = self.__capacity__
        new_slots: t.CPtr = self.__mbuddy__.alloc(new_cap * _HT_SLOT_SIZE)
        if new_slots == None:
            return
        string.memset(new_slots, 0, new_cap * _HT_SLOT_SIZE)
        self.__slots__ = new_slots
        self.__capacity__ = new_cap
        self.__count__ = 0
        self.__tombstones__ = 0
        for i in range(old_cap):
            old_h_ptr: uint64* = t.CPtr(uint64(old_slots) + i * _HT_SLOT_SIZE)
            h: uint64 = old_h_ptr[0]
            if h != _HT_EMPTY and h != _HT_TOMBSTONE:
                old_k_ptr: str* = (str*)(uint64(old_slots) + i * _HT_SLOT_SIZE + 8)
                # Same as _slot_value: the value slot is annotated as "pointer to pointer" (i64**) so the element is a pointer
                old_v_ptr: uint64** = (uint64**)(uint64(old_slots) + i * _HT_SLOT_SIZE + 16)
                k: str = old_k_ptr[0]
                v: t.CPtr = old_v_ptr[0]
                idx: sizet = self._find_slot(k, h)
                self._set_slot(idx, h, k, v)
                self.__count__ += 1

    def __setitem__(self, key: str, val: t.CNeedPtr):
        """Insert or update key-value pair. val automatically decays to pointer (t.CNeedPtr: no conversion if already pointer)."""
        h: uint64 = _fnv1a_hash(key)
        idx: sizet = self._find_slot(key, h)
        if idx >= self.__capacity__:
            return
        sh: uint64 = self._slot_hash(idx)
        if sh == _HT_EMPTY or sh == _HT_TOMBSTONE:
            if sh == _HT_TOMBSTONE:
                self.__tombstones__ -= 1
            self._set_slot(idx, h, key, val)
            self.__count__ += 1
            # Resize when load factor > 0.75
            load: sizet = self.__count__ + self.__tombstones__
            if load * 4 >= self.__capacity__ * 3:
                self._resize(self.__capacity__ * 2)
        else:
            self._set_slot(idx, h, key, val)

    def __getitem__(self, key: str) -> t.CPtr:
        """Lookup key and return value pointer. Returns None if key not found."""
        h: uint64 = _fnv1a_hash(key)
        idx: sizet = self._find_slot(key, h)
        if idx >= self.__capacity__:
            return None
        sh: uint64 = self._slot_hash(idx)
        if sh == _HT_EMPTY or sh == _HT_TOMBSTONE:
            return None
        v: t.CPtr = self._slot_value(idx)
        return v

    def __contains__(self, key: str) -> int:
        """Check whether key exists in table."""
        h: uint64 = _fnv1a_hash(key)
        idx: sizet = self._find_slot(key, h)
        if idx >= self.__capacity__:
            return 0
        sh: uint64 = self._slot_hash(idx)
        if sh == _HT_EMPTY or sh == _HT_TOMBSTONE:
            return 0
        return 1

    def __delitem__(self, key: str) -> int:
        """Delete entry by key. Return 1 on success, 0 if key not present."""
        h: uint64 = _fnv1a_hash(key)
        idx: sizet = self._find_slot(key, h)
        if idx >= self.__capacity__:
            return 0
        sh: uint64 = self._slot_hash(idx)
        if sh == _HT_EMPTY or sh == _HT_TOMBSTONE:
            return 0
        h_ptr: uint64* = t.CPtr(uint64(self.__slots__) + idx * _HT_SLOT_SIZE)
        h_ptr[0] = _HT_TOMBSTONE
        self.__count__ -= 1
        self.__tombstones__ += 1
        return 1

    def __len__(self) -> sizet:
        return self.__count__

    def __iter__(self) -> HashTable*:
        self.__iter_index__ = 0
        return self

    def __next__(self) -> str:
        while self.__iter_index__ < self.__capacity__:
            idx: sizet = self.__iter_index__
            self.__iter_index__ = idx + 1
            sh: uint64 = self._slot_hash(idx)
            if sh != _HT_EMPTY and sh != _HT_TOMBSTONE:
                return self._slot_key(idx)
        raise StopIteration

    def get(self, key: str, default: t.CPtr) -> t.CPtr:
        """Return value pointer; return default pointer if key not found."""
        h: uint64 = _fnv1a_hash(key)
        idx: sizet = self._find_slot(key, h)
        if idx >= self.__capacity__:
            return default
        sh: uint64 = self._slot_hash(idx)
        if sh == _HT_EMPTY or sh == _HT_TOMBSTONE:
            return default
        return self._slot_value(idx)

    def set_int(self, key: str, val: int):
        """Store integer value. Value is copied into mbuddy-managed storage."""
        storage: t.CPtr = self.__mbuddy__.alloc(8)
        int_ptr: int* = t.CPtr(uint64(storage))
        int_ptr[0] = val
        self[key] = storage

    def set_str(self, key: str, val: str):
        """Store string value. Directly store the string pointer."""
        self[key] = val