package tool

// TransformMap transforms map of type IN into a map of type OUT.
// The main point of this function is to optimize memory usage. It allocates a single slice for
// the elements in values of the map rather than a slice per key.
func TransformMap[K comparable, E, V1, V2 any, IN ~map[K]V1, OUT ~map[K]V2](
	in IN,
	v1Len func(V1) int,
	appendV1ToSlice func(V1, []E) []E,
	elemToV2 func([]E) V2,
) OUT {

	if len(in) == 0 {
		return nil
	}

	a := Accumulator[E, V1]{
		Append: appendV1ToSlice,
	}

	result := make(OUT, len(in))
	for _, v := range in {
		a.Grow(v1Len(v))
	}
	for k, v := range in {
		result[k] = elemToV2(a.Clone(v))
	}
	return result
}

func SliceToMap[K comparable, E, V1, V2 any, IN ~[]V1, OUT ~map[K]V2](
	in IN,
	v1Len func(V1) int,
	appendV1ToSlice func(V1, []E) []E,
	elemToV2 func(V1, []E) (K, V2),
) OUT {

	if len(in) == 0 {
		return nil
	}

	a := Accumulator[E, V1]{
		Append: appendV1ToSlice,
	}

	result := make(OUT, len(in))
	for _, v := range in {
		a.Grow(v1Len(v))
	}
	for _, v1 := range in {
		key, v2 := elemToV2(v1, a.Clone(v1))
		result[key] = v2
	}
	return result
}

func MapToSlice[K comparable, E, V1, V2 any, IN ~map[K]V1, OUT ~[]V2](
	in IN,
	v1Len func(V1) int,
	appendV1ToSlice func(V1, []E) []E,
	elemToV2 func(K, []E) V2,
) OUT {

	if len(in) == 0 {
		return nil
	}

	a := Accumulator[E, V1]{
		Append: appendV1ToSlice,
	}

	result := make(OUT, 0, len(in))
	for _, v := range in {
		a.Grow(v1Len(v))
	}
	for k, v1 := range in {
		result = append(result, elemToV2(k, a.Clone(v1)))
	}
	return result
}

func StringMapToSliceWithBytes[V2 any, IN ~map[string][]string, OUT ~[]V2](
	from IN,
	elemToV2 func(string, [][]byte) V2,
) OUT {
	a := Accumulator[byte, string]{
		Append: func(v string, sink []byte) []byte {
			return append(sink, v...)
		},
	}
	return MapToSlice[string, []byte, []string, V2, IN, OUT](
		from,
		func(vals []string) int {
			for _, v := range vals {
				a.Grow(len(v))
			}
			return len(vals)
		},
		func(vals []string, sink [][]byte) [][]byte {
			for _, val := range vals {
				sink = append(sink, a.Clone(val)) // add the copied slice to the sink
			}
			return sink
		},
		elemToV2,
	)
}

// Accumulator allows to reduce the number of allocations by packing multiple slices into a single big slice.
type Accumulator[E, V any] struct {
	Append func(V, []E) []E
	data   []E
	size   int
}

func (a *Accumulator[E, V]) Grow(size int) {
	if a.data != nil {
		panic("Grow called after Clone")
	}
	a.size += size
}

func (a *Accumulator[E, V]) Clone(v V) []E {
	if a.data == nil {
		// Allocate backing array for all elements in one go on first use.
		a.data = make([]E, 0, a.size)
	}
	a.data = a.Append(v, a.data)
	ld := len(a.data)
	result := a.data[:ld:ld] // Set capacity to length to protect against potential append overwriting the next value
	a.data = a.data[ld:]

	return result
}
