package dataext

import (
	"fmt"
	"iter"
)

// OrderedMap is like a normal map[TKey, TVal] - but its elements stay in order
// NOT THREADSAFE !!!
type OrderedMap[TKey comparable, TVal any] struct {
	m map[TKey]*TVal
	a []TKey
}

func NewOrderedMap[TKey comparable, TVal any](caps ...int) *OrderedMap[TKey, TVal] {
	if len(caps) == 0 {
		return &OrderedMap[TKey, TVal]{
			m: make(map[TKey]*TVal),
			a: make([]TKey, 0),
		}
	}

	omcap := 0
	for _, v := range caps {
		omcap += v
	}

	return &OrderedMap[TKey, TVal]{
		m: make(map[TKey]*TVal, omcap),
		a: make([]TKey, 0, omcap),
	}
}

func (o *OrderedMap[TKey, TVal]) Get(key TKey) (TVal, bool) {
	v, ok := o.m[key]
	if ok {
		return *v, ok
	}

	return *new(TVal), false
}

func (o *OrderedMap[TKey, TVal]) GetOrNil(key TKey) *TVal {
	v, ok := o.m[key]
	if ok {
		return v
	}

	return nil
}

func (o *OrderedMap[TKey, TVal]) GetOrDefault(key TKey, defaultVal TVal) TVal {
	v, ok := o.m[key]
	if ok {
		return *v
	}

	return defaultVal
}

// Add adds the new value to the map
// At the end of the map-ordering (even if key already exists, its then "moved" to the end)
// returns true if key already existed
func (o *OrderedMap[TKey, TVal]) Add(key TKey, val TVal) bool {
	if _, ok := o.m[key]; ok {
		o.remFromArray(key)
		o.m[key] = &val
		o.a = append(o.a, key)
		return true
	}

	o.m[key] = &val
	o.a = append(o.a, key)
	return false
}

// AddOrReplace adds the new value to the map
// Normally at the end of the map, but if teh key already exists, its only replaced
// returns true if key already existed
func (o *OrderedMap[TKey, TVal]) AddOrReplace(key TKey, val TVal) bool {
	if _, ok := o.m[key]; ok {
		o.m[key] = &val
		return true
	}

	o.m[key] = &val
	o.a = append(o.a, key)
	return false
}

func (o *OrderedMap[TKey, TVal]) Remove(key TKey) bool {
	if _, ok := o.m[key]; ok {
		o.remFromArray(key)
		delete(o.m, key)
		return true
	}

	return false
}

func (o *OrderedMap[TKey, TVal]) Iterate() iter.Seq2[TKey, TVal] {
	return func(yield func(TKey, TVal) bool) {
		for _, v := range o.a {
			if !yield(v, *o.m[v]) {
				return
			}
		}
	}
}

func (o *OrderedMap[TKey, TVal]) IterateValues() iter.Seq[TVal] {
	return func(yield func(TVal) bool) {
		for _, v := range o.a {
			if !yield(*o.m[v]) {
				return
			}
		}
	}
}

func (o *OrderedMap[TKey, TVal]) IterateKeys() iter.Seq[TKey] {
	return func(yield func(TKey) bool) {
		for _, v := range o.a {
			if !yield(v) {
				return
			}
		}
	}
}

func (o *OrderedMap[TKey, TVal]) Array() []TVal {
	res := make([]TVal, len(o.a))
	for i, v := range o.a {
		res[i] = *o.m[v]
	}

	return res
}

func (o *OrderedMap[TKey, TVal]) Keys() []TKey {
	res := make([]TKey, len(o.a))
	for i, v := range o.a {
		res[i] = v
	}

	return res
}

func (o *OrderedMap[TKey, TVal]) Clear() {
	mapCap := max(len(o.m), cap(o.a))
	o.m = make(map[TKey]*TVal, mapCap)
	o.a = make([]TKey, 0, mapCap)
}

func (o *OrderedMap[TKey, TVal]) Size() int {
	return len(o.a)
}

func (o *OrderedMap[TKey, TVal]) Capacity() int {
	return min(cap(o.a), len(o.m))
}

func (o *OrderedMap[TKey, TVal]) Contains(key TKey) bool {
	_, ok := o.m[key]
	return ok
}

func (o *OrderedMap[TKey, TVal]) IndexOf(key TKey) int {
	for i, v := range o.a {
		if v == key {
			return i
		}
	}

	return -1
}

func (o *OrderedMap[TKey, TVal]) remFromArray(key TKey) {
	for i, v := range o.a {
		if v == key {
			o.a = append(o.a[:i], o.a[i+1:]...)
			return
		}
	}

	panic(fmt.Sprintf("Failed to remove key from OrderedMap -- key '%v' not found", key))
}