package dict

import (
	"hash/fnv"
	"math"
	"runtime"
)

// ShardSafeDict 定义分片且并发安全的数据结构
type ShardSafeDict[K comparable, V any] struct {
	count  int
	shards []*SafeDict[K, V]
}

/*
NewShardSafeDict 创建一个分片的键值对缓存
  - 分片数量：runtime.NumCPU() * 32
*/
func NewShardSafeDict[K comparable, V any]() *ShardSafeDict[K, V] {
	count := runtime.NumCPU() * 32
	shards := make([]*SafeDict[K, V], count)
	for i := range shards {
		shards[i] = NewSafeDict[K, V]()
	}

	return &ShardSafeDict[K, V]{count: count, shards: shards}
}

// Set 设置键值对
func (own *ShardSafeDict[K, V]) Set(key K, val V) {
	dict := own.getShard(key)
	dict.Set(key, val)
}

// Update 更新字典
func (own *ShardSafeDict[K, V]) Update(items map[K]V) {
	for key, val := range items {
		dict := own.getShard(key)
		dict.Set(key, val)
	}
}

// Delete 删除键值对
func (own *ShardSafeDict[K, V]) Delete(keys ...K) {
	for _, key := range keys {
		dict := own.getShard(key)
		dict.Delete(key)
	}
}

// HasAll 判断所有指定键值对是否存在
func (own *ShardSafeDict[K, V]) HasAll(keys ...K) bool {
	for _, key := range keys {
		dict := own.getShard(key)
		if ok := dict.HasAny(key); !ok {
			return false
		}
	}

	return true
}

// HasAny 判断任一指定元素是存在
func (own *ShardSafeDict[K, V]) HasAny(keys ...K) bool {
	for _, key := range keys {
		dict := own.getShard(key)
		if ok := dict.HasAny(key); ok {
			return true
		}
	}

	return false
}

// Get 查询键值对
func (own *ShardSafeDict[K, V]) Get(key K) (V, bool) {
	dict := own.getShard(key)
	return dict.Get(key)
}

// GetLength 获取字典长度
func (own *ShardSafeDict[K, V]) GetLength() (length int) {
	for _, dict := range own.shards {
		length += dict.GetLength()
	}

	return
}

// GetKeys 获取字典中所有的键
func (own *ShardSafeDict[K, V]) GetKeys() []K {
	keys := make([]K, 0)
	for _, dict := range own.shards {
		keys = append(keys, dict.GetKeys()...)
	}

	return keys
}

// GetValues 获取字典中所有的值
func (own *ShardSafeDict[K, V]) GetValues() []V {
	values := make([]V, 0)
	for _, dict := range own.shards {
		values = append(values, dict.GetValues()...)
	}

	return values
}

// GetMap 获取Map
func (own *ShardSafeDict[K, V]) GetMap() map[K]V {
	items := make(map[K]V)
	for _, dict := range own.shards {
		for k, v := range dict.GetMap() {
			items[k] = v
		}
	}

	return items
}

// Range 遍历字典
func (own *ShardSafeDict[K, V]) Range(fn func(key K, val V)) {
	for _, dict := range own.shards {
		dict.Range(fn)
	}
}

// Filter 从字典中过滤出符合要求的键值对
func (own *ShardSafeDict[K, V]) Filter(fn func(key K, val V) bool) map[K]V {
	items := make(map[K]V)
	for _, dict := range own.shards {
		for k, v := range dict.Filter(fn) {
			items[k] = v
		}
	}

	return items
}

// Transform 遍历并更新每个元素
func (own *ShardSafeDict[K, V]) Transform(fn func(key K, val V) V) {
	for _, dict := range own.shards {
		dict.Transform(fn)
	}
}

// Clear 清空字典
func (own *ShardSafeDict[K, V]) Clear() {
	for _, dict := range own.shards {
		dict.Clear()
	}
}

func (own *ShardSafeDict[K, V]) getShard(key K) *SafeDict[K, V] {
	var hash uint64
	switch k := any(key).(type) {
	case string:
		hashFNV := fnv.New64a()
		_, _ = hashFNV.Write([]byte(k))
		hash = hashFNV.Sum64()
	case int:
		hash = uint64(k) ^ (uint64(k) >> 32)
	case int8:
		hash = uint64(k) ^ (uint64(k) >> 32)
	case int16:
		hash = uint64(k) ^ (uint64(k) >> 32)
	case int32:
		hash = uint64(k) ^ (uint64(k) >> 32)
	case int64:
		hash = uint64(k) ^ (uint64(k) >> 32)
	case uint:
		hash = uint64(k)
	case uint8:
		hash = uint64(k)
	case uint16:
		hash = uint64(k)
	case uint32:
		hash = uint64(k)
	case uint64:
		hash = k
	case float32:
		hash = math.Float64bits(float64(k))
	case float64:
		hash = math.Float64bits(k)
	}

	return own.shards[hash%uint64(own.count)]
}