package simpleset

import "maps"

type Set[T comparable] struct {
	contents map[T]struct{}
}

func New[T comparable](elements ...T) *Set[T] {
	s := &Set[T]{
		contents: make(map[T]struct{}),
	}
	for _, e := range elements {
		s.Add(e)
	}
	return s
}

func (s *Set[T]) Add(element T) {
	s.contents[element] = struct{}{}
}

func (s *Set[T]) Remove(element T) {
	delete(s.contents, element)
}

func (s *Set[T]) Contains(element T) bool {
	_, exists := s.contents[element]
	return exists
}

func (s *Set[T]) Size() int {
	return len(s.contents)
}

func (s *Set[T]) ToSlice() []T {
	elements := make([]T, 0, len(s.contents))
	for e := range s.contents {
		elements = append(elements, e)
	}
	return elements
}

func (s *Set[T]) Equal(other *Set[T]) bool {
	return maps.Equal(s.contents, other.contents)
}

func (s *Set[T]) Union(other *Set[T]) *Set[T] {
	result := New[T]()
	for e := range s.contents {
		result.Add(e)
	}
	for e := range other.contents {
		result.Add(e)
	}
	return result
}

func (s *Set[T]) Intersection(other *Set[T]) *Set[T] {
	result := New[T]()
	for e := range s.contents {
		if other.Contains(e) {
			result.Add(e)
		}
	}
	return result
}

func (s *Set[T]) Difference(other *Set[T]) *Set[T] {
	result := New[T]()
	for e := range s.contents {
		if !other.Contains(e) {
			result.Add(e)
		}
	}
	return result
}