package container

import "core:intrinsics"
_ :: intrinsics;


Map :: struct($Key, $Value: typeid) where intrinsics.type_is_valid_map_key(Key) {
	hash: Array(int),
	entries: Array(Map_Entry(Key, Value)),
}

Map_Entry :: struct($Key, $Value: typeid) where intrinsics.type_is_valid_map_key(Key) {
	hash:  uintptr,
	next:  int,
	key:   Key,
	value: Value,
}


/*
map_init :: proc{
	map_init_none,
	map_init_cap,
}
map_delete

map_has
map_get
map_get_default
map_get_ptr
map_set
map_remove
map_reserve
map_clear

// Multi Map

multi_map_find_first
multi_map_find_next
multi_map_count
multi_map_get :: proc{
	multi_map_get_array,
	multi_map_get_slice,
};
multi_map_get_as_slice
multi_map_insert
multi_map_remove
multi_map_remove_all

*/

map_init :: proc{map_init_none, map_init_cap};

map_init_none :: proc(m: ^$M/Map($Key, $Value), allocator := context.allocator) {
	m.hash.allocator = allocator;
	m.entries.allocator = allocator;
}

map_init_cap :: proc(m: ^$M/Map($Key, $Value), cap: int, allocator := context.allocator) {
	m.hash.allocator = allocator;
	m.entries.allocator = allocator;
	map_reserve(m, cap);
}

map_delete :: proc(m: $M/Map($Key, $Value)) {
	array_delete(m.hash);
	array_delete(m.entries);
}


map_has :: proc(m: $M/Map($Key, $Value), key: Key) -> bool {
	return _map_find_or_fail(m, key) >= 0;
}

map_get :: proc(m: $M/Map($Key, $Value), key: Key) -> (res: Value, ok: bool) #optional_ok {
	i := _map_find_or_fail(m, key);
	if i < 0 {
		return {}, false;
	}
	return array_get(m.entries, i).value, true;
}

map_get_default :: proc(m: $M/Map($Key, $Value), key: Key, default: Value) -> (res: Value, ok: bool) #optional_ok {
	i := _map_find_or_fail(m, key);
	if i < 0 {
		return default, false;
	}
	return array_get(m.entries, i).value, true;
}

map_get_ptr :: proc(m: $M/Map($Key, $Value), key: Key) -> ^Value {
	i := _map_find_or_fail(m, key);
	if i < 0 {
		return nil;
	}
	return array_get_ptr(m.entries, i).value;
}

map_set :: proc(m: ^$M/Map($Key, $Value), key: Key, value: Value) {
	if array_len(m.hash) == 0 {
		_map_grow(m);
	}

	i := _map_find_or_make(m, key);
	array_get_ptr(m.entries, i).value = value;
	if _map_full(m^) {
		_map_grow(m);
	}
}

map_remove :: proc(m: ^$M/Map($Key, $Value), key: Key) {
	fr := _map_find_key(m^, key);
	if fr.entry_index >= 0 {
		_map_erase(m, fr);
	}
}


map_reserve :: proc(m: ^$M/Map($Key, $Value), new_size: int) {
	nm: M;
	map_init(&nm, m.hash.allocator);
	array_resize(&nm.hash, new_size);
	array_reserve(&nm.entries, array_len(m.entries));

	for i in 0..<new_size {
		array_set(&nm.hash, i, -1);
	}
	for i in 0..<array_len(m.entries) {
		e := array_get(m.entries, i);
		multi_map_insert(&nm, e.key, e.value);
	}

	map_delete(m^);
	m^ = nm;
}

map_clear :: proc(m: ^$M/Map($Key, $Value)) {
	array_clear(&m.hash);
	array_clear(&m.entries);
}



multi_map_find_first :: proc(m: $M/Map($Key, $Value), key: Key) -> ^Map_Entry(Key, Value) {
	i := _map_find_or_fail(m, key);
	if i < 0 {
		return nil;
	}
	return array_get_ptr(m.entries, i);
}

multi_map_find_next :: proc(m: $M/Map($Key, $Value), e: ^Map_Entry(Key, Value)) -> ^Map_Entry(Key, Value) {
	i := e.next;
	for i >= 0 {
		it := array_get_ptr(m.entries, i);
		if it.hash == e.hash && it.key == e.key {
			return it;
		}
		i = it.next;
	}
	return nil;
}

multi_map_count :: proc(m: $M/Map($Key, $Value), key: Key) -> int {
	n := 0;
	e := multi_map_find_first(m, key);
	for e != nil {
		n += 1;
		e = multi_map_find_next(m, e);
	}
	return n;
}

multi_map_get :: proc{multi_map_get_array, multi_map_get_slice};

multi_map_get_array :: proc(m: $M/Map($Key, $Value), key: Key, items: ^Array(Value)) {
	if items == nil {
		return;
	}
	e := multi_map_find_first(m, key);
	for e != nil {
		array_append(items, e.value);
		e = multi_map_find_next(m, e);
	}
}

multi_map_get_slice :: proc(m: $M/Map($Key, $Value), key: Key, items: []Value) {
	e := multi_map_find_first(m, key);
	i := 0;
	for e != nil && i < len(items) {
		items[i] = e.value;
		i += 1;
		e = multi_map_find_next(m, e);
	}
}

multi_map_get_as_slice :: proc(m: $M/Map($Key, $Value), key: Key) -> []Value {
	items: Array(Value);
	array_init(&items, 0);

	e := multi_map_find_first(m, key);
	for e != nil {
		array_append(&items, e.value);
		e = multi_map_find_next(m, e);
	}

	return array_slice(items);
}


multi_map_insert :: proc(m: ^$M/Map($Key, $Value), key: Key, value: Value) {
	if array_len(m.hash) == 0 {
		_map_grow(m);
	}

	i := _map_make(m, key);
	array_get_ptr(m.entries, i).value = value;
	if _map_full(m^) {
		_map_grow(m);
	}
}

multi_map_remove :: proc(m: ^$M/Map($Key, $Value), e: ^Map_Entry(Key, Value)) {
	fr := _map_find_entry(m, e);
	if fr.entry_index >= 0 {
		_map_erase(m, fr);
	}
}

multi_map_remove_all :: proc(m: ^$M/Map($Key, $Value), key: Key) {
	for map_exist(m^, key) {
		map_remove(m, key);
	}
}


/// Internal


Map_Find_Result :: struct {
	hash_index:  int,
	entry_prev:  int,
	entry_index: int,
}

_map_add_entry :: proc(m: ^$M/Map($Key, $Value), key: Key) -> int where intrinsics.type_is_valid_map_key(Key) {
	hasher := intrinsics.type_hasher_proc(Key);

	e: Map_Entry(Key, Value);
	e.key = key;
	e.hash = hasher(&e.key, 0);
	e.next = -1;
	idx := array_len(m.entries);
	array_push(&m.entries, e);
	return idx;
}

_map_erase :: proc(m: ^$M/Map, fr: Map_Find_Result) {
	if fr.entry_prev < 0 {
		array_set(&m.hash, fr.hash_index, array_get(m.entries, fr.entry_index).next);
	} else {
		array_get_ptr(m.entries, fr.entry_prev).next = array_get(m.entries, fr.entry_index).next;
	}

	if fr.entry_index == array_len(m.entries)-1 {
		array_pop_back(&m.entries);
		return;
	}

	array_set(&m.entries, fr.entry_index, array_get(m.entries, array_len(m.entries)-1));
	last := _map_find_key(m^, array_get(m.entries, fr.entry_index).key);

	if last.entry_prev < 0 {
		array_get_ptr(m.entries, last.entry_prev).next = fr.entry_index;
	} else {
		array_set(&m.hash, last.hash_index, fr.entry_index);
	}
}


_map_find_key :: proc(m: $M/Map($Key, $Value), key: Key) -> Map_Find_Result where intrinsics.type_is_valid_map_key(Key) {
	fr: Map_Find_Result;
	fr.hash_index = -1;
	fr.entry_prev = -1;
	fr.entry_index = -1;

	if array_len(m.hash) == 0 {
		return fr;
	}

	hasher := intrinsics.type_hasher_proc(Key);

	key := key;
	hash := hasher(&key, 0);

	fr.hash_index = int(hash % uintptr(array_len(m.hash)));
	fr.entry_index = array_get(m.hash, fr.hash_index);
	for fr.entry_index >= 0 {
		it := array_get_ptr(m.entries, fr.entry_index);
		if it.hash == hash && it.key == key {
			return fr;
		}
		fr.entry_prev = fr.entry_index;
		fr.entry_index = it.next;
	}
	return fr;
}

_map_find_entry :: proc(m: ^$M/Map($Key, $Value), e: ^Map_Entry(Key, Value)) -> Map_Find_Result {
	fr: Map_Find_Result;
	fr.hash_index = -1;
	fr.entry_prev = -1;
	fr.entry_index = -1;

	if array_len(m.hash) == 0 {
		return fr;
	}

	fr.hash_index = int(e.hash % uintptr(array_len(m.hash)));
	fr.entry_index = array_get(m.hash, fr.hash_index);
	for fr.entry_index >= 0 {
		it := array_get_ptr(m.entries, fr.entry_index);
		if it == e {
			return fr;
		}
		fr.entry_prev = fr.entry_index;
		fr.entry_index = it.next;
	}
	return fr;
}

_map_find_or_fail :: proc(m: $M/Map($Key, $Value), key: Key) -> int {
	return _map_find_key(m, key).entry_index;
}
_map_find_or_make :: proc(m: ^$M/Map($Key, $Value), key: Key) -> int {
	fr := _map_find_key(m^, key);
	if fr.entry_index >= 0 {
		return fr.entry_index;
	}

	i := _map_add_entry(m, key);
	if fr.entry_prev < 0 {
		array_set(&m.hash, fr.hash_index, i);
	} else {
		array_get_ptr(m.entries, fr.entry_prev).next = i;
	}
	return i;
}


_map_make :: proc(m: ^$M/Map($Key, $Value), key: Key) -> int {
	fr := _map_find_key(m^, key);
	i := _map_add_entry(m, key);

	if fr.entry_prev < 0 {
		array_set(&m.hash, fr.hash_index, i);
	} else {
		array_get_ptr(m.entries, fr.entry_prev).next = i;
	}

	array_get_ptr(m.entries, i).next = fr.entry_index;

	return i;
}


_map_full :: proc(m: $M/Map($Key, $Value)) -> bool {
	// TODO(bill): Determine good max load factor
	return array_len(m.entries) >= (array_len(m.hash) / 4)*3;
}

_map_grow :: proc(m: ^$M/Map($Key, $Value)) {
	new_size := array_len(m.entries) * 4 + 7; // TODO(bill): Determine good grow rate
	map_reserve(m, new_size);
}