package dataext

import "sync"

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{}}
}

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
}

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
}

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
}

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
	}
}

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
	}
}

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
	}
}

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
}

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
}

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
}

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
}