package dataext

import "sync"

// SyncMap is a thread-safe map implementation for generic key-value pairs.
// All functions aresafe to be called in parallel.
type SyncMap[TKey comparable, TData any] struct {
	data map[TKey]TData
	lock sync.Mutex
}

func NewSyncMap[TKey comparable, TData any]() *SyncMap[TKey, TData] {
	return &SyncMap[TKey, TData]{data: make(map[TKey]TData), lock: sync.Mutex{}}
}

// Set sets the value for the provided key
func (s *SyncMap[TKey, TData]) Set(key TKey, data TData) {
	s.lock.Lock()
	defer s.lock.Unlock()

	if s.data == nil {
		s.data = make(map[TKey]TData)
	}

	s.data[key] = data
}

// SetIfNotContains sets the value for the provided key if it does not already exist
func (s *SyncMap[TKey, TData]) SetIfNotContains(key TKey, data TData) bool {
	s.lock.Lock()
	defer s.lock.Unlock()

	if s.data == nil {
		s.data = make(map[TKey]TData)
	}

	if _, existsInPreState := s.data[key]; existsInPreState {
		return false
	}

	s.data[key] = data

	return true
}

// SetIfNotContainsFunc sets the value for the provided key using the provided function
func (s *SyncMap[TKey, TData]) SetIfNotContainsFunc(key TKey, data func() TData) bool {
	s.lock.Lock()
	defer s.lock.Unlock()

	if s.data == nil {
		s.data = make(map[TKey]TData)
	}

	if _, existsInPreState := s.data[key]; existsInPreState {
		return false
	}

	s.data[key] = data()

	return true
}

// Get retrieves the value for the provided key
func (s *SyncMap[TKey, TData]) Get(key TKey) (TData, bool) {
	s.lock.Lock()
	defer s.lock.Unlock()

	if s.data == nil {
		s.data = make(map[TKey]TData)
	}

	if v, ok := s.data[key]; ok {
		return v, true
	} else {
		return *new(TData), false
	}
}

// GetAndSetIfNotContains returns the existing value if the key exists.
// Otherwise, it sets the provided value and returns it.
func (s *SyncMap[TKey, TData]) GetAndSetIfNotContains(key TKey, data TData) TData {
	s.lock.Lock()
	defer s.lock.Unlock()

	if s.data == nil {
		s.data = make(map[TKey]TData)
	}

	if v, ok := s.data[key]; ok {
		return v
	} else {
		s.data[key] = data
		return data
	}
}

// GetAndSetIfNotContainsFunc returns the existing value if the key exists.
// Otherwise, it calls the provided function to generate the value, sets it, and returns it.
func (s *SyncMap[TKey, TData]) GetAndSetIfNotContainsFunc(key TKey, data func() TData) TData {
	s.lock.Lock()
	defer s.lock.Unlock()

	if s.data == nil {
		s.data = make(map[TKey]TData)
	}

	if v, ok := s.data[key]; ok {
		return v
	} else {
		dataObj := data()
		s.data[key] = dataObj
		return dataObj
	}
}

// Delete removes the entry with the provided key and returns true if the key existed before.
func (s *SyncMap[TKey, TData]) Delete(key TKey) bool {
	s.lock.Lock()
	defer s.lock.Unlock()

	if s.data == nil {
		s.data = make(map[TKey]TData)
	}

	_, ok := s.data[key]

	delete(s.data, key)

	return ok
}

// DeleteIf deletes all entries that match the provided function and returns the number of removed entries.
func (s *SyncMap[TKey, TData]) DeleteIf(fn func(key TKey, data TData) bool) int {
	s.lock.Lock()
	defer s.lock.Unlock()

	if s.data == nil {
		s.data = make(map[TKey]TData)
	}

	rm := 0
	for k, v := range s.data {
		if fn(k, v) {
			delete(s.data, k)
			rm++
		}
	}

	return rm
}

// UpdateIfExists updates the value if the key exists, otherwise it does nothing.
func (s *SyncMap[TKey, TData]) UpdateIfExists(key TKey, fn func(data TData) TData) bool {
	s.lock.Lock()
	defer s.lock.Unlock()

	if s.data == nil {
		s.data = make(map[TKey]TData)
	}

	if v, ok := s.data[key]; ok {
		s.data[key] = fn(v)
		return true
	} else {
		return false
	}
}

// UpdateOrInsert updates the value if the key exists, otherwise it inserts the provided `insertValue`.
func (s *SyncMap[TKey, TData]) UpdateOrInsert(key TKey, fn func(data TData) TData, insertValue TData) bool {
	s.lock.Lock()
	defer s.lock.Unlock()

	if s.data == nil {
		s.data = make(map[TKey]TData)
	}

	if v, ok := s.data[key]; ok {
		s.data[key] = fn(v)
		return true
	} else {
		s.data[key] = insertValue
		return false
	}
}

// Clear removes all entries from the map.
func (s *SyncMap[TKey, TData]) Clear() {
	s.lock.Lock()
	defer s.lock.Unlock()

	s.data = make(map[TKey]TData)
}

// Contains checks if the map contains the provided key.
func (s *SyncMap[TKey, TData]) Contains(key TKey) bool {
	s.lock.Lock()
	defer s.lock.Unlock()

	if s.data == nil {
		s.data = make(map[TKey]TData)
	}

	_, ok := s.data[key]

	return ok
}

// GetAllKeys returns a copy (!) of all keys in the map.
func (s *SyncMap[TKey, TData]) GetAllKeys() []TKey {
	s.lock.Lock()
	defer s.lock.Unlock()

	if s.data == nil {
		s.data = make(map[TKey]TData)
	}

	r := make([]TKey, 0, len(s.data))

	for k := range s.data {
		r = append(r, k)
	}

	return r
}

// GetAllValues returns a copy (!) of all values in the map.
func (s *SyncMap[TKey, TData]) GetAllValues() []TData {
	s.lock.Lock()
	defer s.lock.Unlock()

	if s.data == nil {
		s.data = make(map[TKey]TData)
	}

	r := make([]TData, 0, len(s.data))

	for _, v := range s.data {
		r = append(r, v)
	}

	return r
}

// Count returns the number of entries in the map.
func (s *SyncMap[TKey, TData]) Count() int {
	s.lock.Lock()
	defer s.lock.Unlock()

	if s.data == nil {
		s.data = make(map[TKey]TData)
	}

	return len(s.data)
}