/*
 * SPDX-FileCopyrightText: © Hypermode Inc. <hello@hypermode.com>
 * SPDX-License-Identifier: Apache-2.0
 */

package ristretto

import (
	
	
)

type updateFn[ any] func(cur, prev ) bool

// TODO: Do we need this to be a separate struct from Item?
type storeItem[ any] struct {
	key        uint64
	conflict   uint64
	value      
	expiration time.Time
}

// store is the interface fulfilled by all hash map implementations in this
// file. Some hash map implementations are better suited for certain data
// distributions than others, so this allows us to abstract that out for use
// in Ristretto.
//
// Every store is safe for concurrent usage.
type store[ any] interface {
	// Get returns the value associated with the key parameter.
	Get(uint64, uint64) (, bool)
	// Expiration returns the expiration time for this key.
	Expiration(uint64) time.Time
	// Set adds the key-value pair to the Map or updates the value if it's
	// already present. The key-value pair is passed as a pointer to an
	// item object.
	Set(*Item[])
	// Del deletes the key-value pair from the Map.
	Del(uint64, uint64) (uint64, )
	// Update attempts to update the key with a new value and returns true if
	// successful.
	Update(*Item[]) (, bool)
	// Cleanup removes items that have an expired TTL.
	Cleanup(policy *defaultPolicy[], onEvict func(item *Item[]))
	// Clear clears all contents of the store.
	Clear(onEvict func(item *Item[]))
	SetShouldUpdateFn(f updateFn[])
}

// newStore returns the default store implementation.
func newStore[ any]() store[] {
	return newShardedMap[]()
}

const numShards uint64 = 256

type shardedMap[ any] struct {
	shards    []*lockedMap[]
	expiryMap *expirationMap[]
}

func newShardedMap[ any]() *shardedMap[] {
	 := &shardedMap[]{
		shards:    make([]*lockedMap[], int(numShards)),
		expiryMap: newExpirationMap[](),
	}
	for  := range .shards {
		.shards[] = newLockedMap[](.expiryMap)
	}
	return 
}

func ( *shardedMap[]) ( updateFn[]) {
	for  := range .shards {
		.shards[].setShouldUpdateFn()
	}
}

func ( *shardedMap[]) (,  uint64) (, bool) {
	return .shards[%numShards].get(, )
}

func ( *shardedMap[]) ( uint64) time.Time {
	return .shards[%numShards].Expiration()
}

func ( *shardedMap[]) ( *Item[]) {
	if  == nil {
		// If item is nil make this Set a no-op.
		return
	}

	.shards[.Key%numShards].Set()
}

func ( *shardedMap[]) (,  uint64) (uint64, ) {
	return .shards[%numShards].Del(, )
}

func ( *shardedMap[]) ( *Item[]) (, bool) {
	return .shards[.Key%numShards].Update()
}

func ( *shardedMap[]) ( *defaultPolicy[],  func( *Item[])) {
	.expiryMap.cleanup(, , )
}

func ( *shardedMap[]) ( func( *Item[])) {
	for  := uint64(0);  < numShards; ++ {
		.shards[].Clear()
	}
	.expiryMap.clear()
}

type lockedMap[ any] struct {
	sync.RWMutex
	data         map[uint64]storeItem[]
	em           *expirationMap[]
	shouldUpdate updateFn[]
}

func newLockedMap[ any]( *expirationMap[]) *lockedMap[] {
	return &lockedMap[]{
		data: make(map[uint64]storeItem[]),
		em:   ,
		shouldUpdate: func(,  ) bool {
			return true
		},
	}
}

func ( *lockedMap[]) ( updateFn[]) {
	.shouldUpdate = 
}

func ( *lockedMap[]) (,  uint64) (, bool) {
	.RLock()
	,  := .data[]
	.RUnlock()
	if ! {
		return zeroValue[](), false
	}
	if  != 0 && ( != .conflict) {
		return zeroValue[](), false
	}

	// Handle expired items.
	if !.expiration.IsZero() && time.Now().After(.expiration) {
		return zeroValue[](), false
	}
	return .value, true
}

func ( *lockedMap[]) ( uint64) time.Time {
	.RLock()
	defer .RUnlock()
	return .data[].expiration
}

func ( *lockedMap[]) ( *Item[]) {
	if  == nil {
		// If the item is nil make this Set a no-op.
		return
	}

	.Lock()
	defer .Unlock()
	,  := .data[.Key]

	if  {
		// The item existed already. We need to check the conflict key and reject the
		// update if they do not match. Only after that the expiration map is updated.
		if .Conflict != 0 && (.Conflict != .conflict) {
			return
		}
		if .shouldUpdate != nil && !.shouldUpdate(.Value, .value) {
			return
		}
		.em.update(.Key, .Conflict, .expiration, .Expiration)
	} else {
		// The value is not in the map already. There's no need to return anything.
		// Simply add the expiration map.
		.em.add(.Key, .Conflict, .Expiration)
	}

	.data[.Key] = storeItem[]{
		key:        .Key,
		conflict:   .Conflict,
		value:      .Value,
		expiration: .Expiration,
	}
}

func ( *lockedMap[]) (,  uint64) (uint64, ) {
	.Lock()
	defer .Unlock()
	,  := .data[]
	if ! {
		return 0, zeroValue[]()
	}
	if  != 0 && ( != .conflict) {
		return 0, zeroValue[]()
	}

	if !.expiration.IsZero() {
		.em.del(, .expiration)
	}

	delete(.data, )
	return .conflict, .value
}

func ( *lockedMap[]) ( *Item[]) (, bool) {
	.Lock()
	defer .Unlock()
	,  := .data[.Key]
	if ! {
		return zeroValue[](), false
	}
	if .Conflict != 0 && (.Conflict != .conflict) {
		return zeroValue[](), false
	}
	if .shouldUpdate != nil && !.shouldUpdate(.Value, .value) {
		return .value, false
	}

	.em.update(.Key, .Conflict, .expiration, .Expiration)
	.data[.Key] = storeItem[]{
		key:        .Key,
		conflict:   .Conflict,
		value:      .Value,
		expiration: .Expiration,
	}

	return .value, true
}

func ( *lockedMap[]) ( func( *Item[])) {
	.Lock()
	defer .Unlock()
	 := &Item[]{}
	if  != nil {
		for ,  := range .data {
			.Key = .key
			.Conflict = .conflict
			.Value = .value
			()
		}
	}
	.data = make(map[uint64]storeItem[])
}