package set

import (
	"maps"
	"sort"

	"golang.org/x/exp/constraints"
)

type Set[T constraints.Ordered] map[T]struct{}

func New[T constraints.Ordered](elements ...T) Set[T] {
	s := Set[T]{}
	for _, e := range elements {
		s.Add(e)
	}
	return s
}

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

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

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

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

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

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

func (s Set[T]) ToSortedSlice() []T {
	elements := s.ToSlice()
	sort.Slice(elements, func(i, j int) bool {
		return elements[i] < elements[j]
	})
	return elements
}

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

func (s Set[T]) Intersection(other Set[T]) Set[T] {
	result := New[T]()
	for e := range s {
		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 {
		if !other.Contains(e) {
			result.Add(e)
		}
	}
	return result
}