package set

import (
	"iter"
	"maps"
	"slices"
)

// Set is a set of elements.
type Set[T comparable] map[T]struct{}

// NewSet creates a new empty set
func NewSet[T comparable]() Set[T] {
	s := make(Set[T])
	return s
}

func NewSetFromSlice[T comparable](items []T) Set[T] {
	s := make(Set[T])
	for _, item := range items {
		s.Insert(item)
	}
	return s
}

// Len returns the number of elements in the set.
func (s Set[T]) Len() int {
	return len(s)
}

// Has returns true if the set contains the specified element.
func (s Set[T]) Has(v T) bool {
	_, ok := s[v]
	return ok
}

// Insert adds the specified element to the set.
func (s Set[T]) Insert(vals ...T) {
	for _, val := range vals {
		s[val] = struct{}{}
	}
}

// Delete removes the specified element from the set.
func (s Set[T]) Delete(v T) {
	delete(s, v)
}

// All returns a iterator over the elements in the set.
func (s Set[T]) All() iter.Seq[T] {
	return maps.Keys(s)
}

// Slice returns a slice containing all the elements in the set.
func (s Set[T]) Slice() []T {
	return slices.Collect(maps.Keys(s))
}

// Merge merge a set.
func (s Set[T]) Merge(other Set[T]) {
	for k := range other {
		s.Insert(k)
	}
}

// Union returns a new set that contains the union of s and other.
func (s Set[T]) Union(other Set[T]) Set[T] {
	result := make(Set[T])
	for k := range s {
		result.Insert(k)
	}
	for k := range other {
		result.Insert(k)
	}
	return result
}

// Intersection returns a new set that contains the intersection of s and other.
func (s Set[T]) Intersection(other Set[T]) Set[T] {
	result := make(Set[T])
	for k := range s {
		if other.Has(k) {
			result.Insert(k)
		}
	}
	return result
}

// Difference returns a new set that contains the elements in s but not in other.
func (s Set[T]) Difference(other Set[T]) Set[T] {
	result := make(Set[T])
	for k := range s {
		if !other.Has(k) {
			result.Insert(k)
		}
	}
	return result
}